[Observability AI Assistant]: Adds several function implementations to the AI Asssistant (#163764)

Lots of small fixes, but mainly new (APM) functions for the
Observability AI Assistant.


f489a310-6ba8-4591-8ac9-54a176e0b58d

---------

Co-authored-by: Coen Warmer <coen.warmer@gmail.com>
Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Co-authored-by: Clint Andrew Hall <clint@clintandrewhall.com>

packages/kbn-es-types/src/search.ts

@@ -162,6 +162,20 @@ export type AggregateOf<
       cardinality: {
         value: number;
       };
+      change_point: {
+        bucket?: {
+          key: string;
+        };
+        type: Record<
+          string,
+          {
+            change_point?: number;
+            r_value?: number;
+            trend?: string;
+            p_value: number;
+          }
+        >;
+      };
       children: {
         doc_count: number;
       } & SubAggregateOf<TAggregationContainer, TDocument>;

packages/kbn-optimizer/limits.yml

@@ -97,7 +97,7 @@ pageLoadAssetSize:
   navigation: 37269
   newsfeed: 42228
   observability: 115443
-  observabilityAIAssistant: 16759
+  observabilityAIAssistant: 25000
   observabilityOnboarding: 19573
   observabilityShared: 52256
   osquery: 107090

x-pack/plugins/apm/common/assistant/constants.ts

@@ -0,0 +1,12 @@
+/*
+ * 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.
+ */
+
+export enum CorrelationsEventType {
+  Transaction = 'transaction',
+  ExitSpan = 'exit_span',
+  Error = 'error',
+}

x-pack/plugins/apm/common/connections.ts

@@ -21,6 +21,7 @@ export interface ServiceNode extends NodeBase {
   serviceName: string;
   agentName: AgentName;
   environment: string;
+  dependencyName?: string;
 }
 
 export interface DependencyNode extends NodeBase {

x-pack/plugins/apm/public/components/assistant_functions/get_apm_correlations.ts

@@ -0,0 +1,118 @@
+/*
+ * 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 '@kbn/observability-ai-assistant-plugin/common/types';
+import { CorrelationsEventType } from '../../../common/assistant/constants';
+import { callApmApi } from '../../services/rest/create_call_apm_api';
+
+export function registerGetApmCorrelationsFunction({
+  registerFunction,
+}: {
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'get_apm_correlations',
+      contexts: ['apm'],
+      description: `Get field values that are more prominent in the foreground set than the 
+      background set. This can be useful in determining what attributes (like 
+      error.message, service.node.name or transaction.name) are contributing to for 
+      instance a higher latency. Another option is a time-based comparison, where you 
+      compare before and after a change point. In KQL, escaping happens with double 
+      quotes, not single quotes. Some characters that need escaping are: ':()\\\/\". 
+      IF you need to filter, make sure the fields are available on the event, and 
+      ALWAYS put a field value in double quotes. Best: event.outcome:\"failure\". 
+      Wrong: event.outcome:'failure'. This is very important! ONLY use this function 
+      if you have something to compare it to.`,
+      descriptionForUser: `Get field values that are more prominent in the foreground set than the 
+      background set. This can be useful in determining what attributes (like 
+      error.message, service.node.name or transaction.name) are contributing to for 
+      instance a higher latency. Another option is a time-based comparison, where you 
+      compare before and after a change point.`,
+      parameters: {
+        type: 'object',
+        properties: {
+          sets: {
+            type: 'array',
+            items: {
+              type: 'object',
+              properties: {
+                label: {
+                  type: 'string',
+                  description:
+                    'A unique, human readable label for the comparison set.',
+                },
+                background: {
+                  description: 'The background data set',
+                  $ref: '#/$defs/set',
+                },
+                foreground: {
+                  description:
+                    'The foreground data set. Needs to be a subset of the background set',
+                  $ref: '#/$defs/set',
+                },
+                event: {
+                  type: 'string',
+                  enum: [
+                    CorrelationsEventType.Error,
+                    CorrelationsEventType.Transaction,
+                    CorrelationsEventType.ExitSpan,
+                  ],
+                },
+              },
+              required: ['background', 'foreground', 'event'],
+            },
+          },
+        },
+        required: ['sets'],
+        $defs: {
+          set: {
+            type: 'object',
+            properties: {
+              'service.name': {
+                type: 'string',
+                description: 'The name of the service',
+              },
+              'service.environment': {
+                type: 'string',
+                description: 'The environment that the service is running in.',
+              },
+              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. Always escape, with double quotes. If no filter should be applied, leave it empty.',
+              },
+              label: {
+                type: 'string',
+                description: 'A unique, human readable label.',
+              },
+            },
+            required: ['service.name', 'start', 'end', 'label'],
+          },
+        },
+      } as const,
+    },
+    async ({ arguments: args }, signal) => {
+      return callApmApi('POST /internal/apm/assistant/get_correlation_values', {
+        signal,
+        params: {
+          body: args,
+        },
+      });
+    }
+  );
+}

x-pack/plugins/apm/public/components/assistant_functions/get_apm_downstream_dependencies.ts

@@ -0,0 +1,65 @@
+/*
+ * 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 '@kbn/observability-ai-assistant-plugin/common/types';
+import { callApmApi } from '../../services/rest/create_call_apm_api';
+
+export function registerGetApmDownstreamDependenciesFunction({
+  registerFunction,
+}: {
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'get_apm_downstream_dependencies',
+      contexts: ['apm'],
+      description: `Get the downstream dependencies (services or uninstrumented backends) for a 
+      service. This allows you to map the dowstream dependency name to a service, by 
+      returning both span.destination.service.resource and service.name. Use this to 
+      drilldown further if needed.`,
+      descriptionForUser: `Get the downstream dependencies (services or uninstrumented backends) for a 
+      service. This allows you to map the dowstream dependency name to a service, by 
+      returning both span.destination.service.resource and service.name. Use this to 
+      drilldown further if needed.`,
+      parameters: {
+        type: 'object',
+        properties: {
+          'service.name': {
+            type: 'string',
+            description: 'The name of the service',
+          },
+          'service.environment': {
+            type: 'string',
+            description: 'The environment that the service is running in',
+          },
+          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`.',
+          },
+        },
+        required: ['service.name', 'start', 'end'],
+      } as const,
+    },
+    async ({ arguments: args }, signal) => {
+      return callApmApi(
+        'GET /internal/apm/assistant/get_downstream_dependencies',
+        {
+          signal,
+          params: {
+            query: args,
+          },
+        }
+      );
+    }
+  );
+}

x-pack/plugins/apm/public/components/assistant_functions/get_apm_error_document.ts

@@ -0,0 +1,56 @@
+/*
+ * 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 '@kbn/observability-ai-assistant-plugin/common/types';
+import { callApmApi } from '../../services/rest/create_call_apm_api';
+
+export function registerGetApmErrorDocumentFunction({
+  registerFunction,
+}: {
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'get_apm_error_document',
+      contexts: ['apm'],
+      description: `Get a sample error document based on its grouping name. This also includes the 
+      stacktrace of the error, which might give you a hint as to what the cause is. 
+      ONLY use this for error events.`,
+      descriptionForUser: `Get a sample error document based on its grouping name. This also includes the 
+      stacktrace of the error, which might give you a hint as to what the cause is.`,
+      parameters: {
+        type: 'object',
+        properties: {
+          'error.grouping_name': {
+            type: 'string',
+            description:
+              'The grouping name of the error. Use the field value returned by get_apm_chart or get_correlation_values.',
+          },
+          start: {
+            type: 'string',
+            description:
+              'The start of the time range, in Elasticsearch date math, lik  e `now`.',
+          },
+          end: {
+            type: 'string',
+            description:
+              'The end of the time range, in Elasticsearch date math, like `now-24h`.',
+          },
+        },
+        required: ['start', 'end', 'error.grouping_name'],
+      } as const,
+    },
+    async ({ arguments: args }, signal) => {
+      return callApmApi('GET /internal/apm/assistant/get_error_document', {
+        signal,
+        params: {
+          query: args,
+        },
+      });
+    }
+  );
+}

x-pack/plugins/apm/public/components/assistant_functions/get_apm_service_summary.ts

@@ -0,0 +1,62 @@
+/*
+ * 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 '@kbn/observability-ai-assistant-plugin/common/types';
+import { callApmApi } from '../../services/rest/create_call_apm_api';
+
+export function registerGetApmServiceSummaryFunction({
+  registerFunction,
+}: {
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'get_apm_service_summary',
+      contexts: ['apm'],
+      description: `Gets a summary of a single service, including: the language, service version, 
+deployments, and the infrastructure that it is running in, for instance on how 
+many pods, and a list of its downstream dependencies. It also returns active 
+alerts and anomalies.`,
+      descriptionForUser: `Gets a summary of a single service, including: the language, service version, 
+deployments, and the infrastructure that it is running in, for instance on how 
+many pods, and a list of its downstream dependencies. It also returns active 
+alerts and anomalies.`,
+      parameters: {
+        type: 'object',
+        properties: {
+          'service.name': {
+            type: 'string',
+            description: 'The name of the service that should be summarized.',
+          },
+          'service.environment': {
+            type: 'string',
+            description: 'The environment that the service is running in',
+          },
+          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`.',
+          },
+        },
+        required: ['service.name', 'start', 'end'],
+      } as const,
+    },
+    async ({ arguments: args }, signal) => {
+      return callApmApi('GET /internal/apm/assistant/get_service_summary', {
+        signal,
+        params: {
+          query: args,
+        },
+      });
+    }
+  );
+}

x-pack/plugins/apm/public/components/assistant_functions/get_apm_timeseries.tsx

@@ -0,0 +1,295 @@
+/*
+ * 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 { EuiFlexGroup, EuiFlexItem, EuiText } from '@elastic/eui';
+import type { RegisterFunctionDefinition } from '@kbn/observability-ai-assistant-plugin/common/types';
+import { groupBy } from 'lodash';
+import React from 'react';
+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 '../shared/charts/helper/timezone';
+import { TimeseriesChart } from '../shared/charts/timeseries_chart';
+import { ChartPointerEventContextProvider } from '../../context/chart_pointer_event/chart_pointer_event_context';
+import { ApmThemeProvider } from '../routing/app_root';
+import { Coordinate, TimeSeries } from '../../../typings/timeseries';
+import {
+  ChartType,
+  getTimeSeriesColor,
+} from '../shared/charts/helper/get_timeseries_color';
+import { LatencyAggregationType } from '../../../common/latency_aggregation_types';
+import {
+  asPercent,
+  asTransactionRate,
+  getDurationFormatter,
+} from '../../../common/utils/formatters';
+import {
+  getMaxY,
+  getResponseTimeTickFormatter,
+} from '../shared/charts/transaction_charts/helper';
+
+export function registerGetApmTimeseriesFunction({
+  registerFunction,
+}: {
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      contexts: ['apm'],
+      name: 'get_apm_timeseries',
+      descriptionForUser: `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: `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. 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!`,
+      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': {
+                  type: 'string',
+                  description: 'The name of the service',
+                },
+                'service.environment': {
+                  type: 'string',
+                  description:
+                    "The environment that the service is running in. If you don't know this, use ENVIRONMENT_ALL.",
+                },
+                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',
+                'service.environment',
+                '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 },
+          },
+        }
+      );
+
+      return response;
+    },
+    ({ arguments: args, response }) => {
+      const groupedSeries = groupBy(response.data, (series) => series.group);
+
+      const {
+        services: { uiSettings },
+      } = useKibana();
+
+      const timeZone = getTimeZone(uiSettings);
+
+      return (
+        <ChartPointerEventContextProvider>
+          <ApmThemeProvider>
+            <EuiFlexGroup direction="column">
+              {Object.values(groupedSeries).map((groupSeries) => {
+                const groupId = groupSeries[0].group;
+
+                const maxY = getMaxY(groupSeries);
+                const latencyFormatter = getDurationFormatter(maxY);
+
+                let yLabelFormat: (value: number) => string;
+
+                const firstStat = groupSeries[0].stat;
+
+                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_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;
+
+                    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_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':
+                        chartType = ChartType.LATENCY_AVG;
+                        break;
+
+                      case 'error_event_rate':
+                        chartType = ChartType.ERROR_OCCURRENCES;
+                        break;
+                    }
+
+                    return {
+                      title: series.id,
+                      type: 'line',
+                      color: getTimeSeriesColor(chartType!).currentPeriodColor,
+                      data: series.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>
+      );
+    }
+  );
+}

x-pack/plugins/apm/public/components/assistant_functions/index.ts

@@ -0,0 +1,132 @@
+/*
+ * 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 { CoreStart } from '@kbn/core/public';
+import {
+  RegisterContextDefinition,
+  RegisterFunctionDefinition,
+} from '@kbn/observability-ai-assistant-plugin/common/types';
+import { ApmPluginStartDeps } from '../../plugin';
+import { createCallApmApi } 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 { registerGetApmServiceSummaryFunction } from './get_apm_service_summary';
+import { registerGetApmTimeseriesFunction } from './get_apm_timeseries';
+
+export function registerAssistantFunctions({
+  pluginsStart,
+  coreStart,
+  registerContext,
+  registerFunction,
+}: {
+  pluginsStart: ApmPluginStartDeps;
+  coreStart: CoreStart;
+  registerFunction: RegisterFunctionDefinition;
+  registerContext: RegisterContextDefinition;
+}) {
+  createCallApmApi(coreStart);
+
+  registerGetApmTimeseriesFunction({
+    registerFunction,
+  });
+
+  registerGetApmErrorDocumentFunction({
+    registerFunction,
+  });
+
+  registerGetApmCorrelationsFunction({
+    registerFunction,
+  });
+
+  registerGetApmDownstreamDependenciesFunction({
+    registerFunction,
+  });
+
+  registerGetApmServiceSummaryFunction({
+    registerFunction,
+  });
+
+  registerContext({
+    name: 'apm',
+    description: `
+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.`,
+  });
+}

x-pack/plugins/apm/public/components/routing/app_root/index.tsx

@@ -127,7 +127,7 @@ function MountApmHeaderActionMenu() {
   );
 }
 
-function ApmThemeProvider({ children }: { children: React.ReactNode }) {
+export function ApmThemeProvider({ children }: { children: React.ReactNode }) {
   const [darkMode] = useUiSetting$<boolean>('theme:darkMode');
 
   return (

x-pack/plugins/apm/public/plugin.ts

@@ -5,10 +5,11 @@
  * 2.0.
  */
 
-import { i18n } from '@kbn/i18n';
-import { from } from 'rxjs';
-import { map } from 'rxjs/operators';
-import { UsageCollectionStart } from '@kbn/usage-collection-plugin/public';
+import type {
+  PluginSetupContract as AlertingPluginPublicSetup,
+  PluginStartContract as AlertingPluginPublicStart,
+} from '@kbn/alerting-plugin/public';
+import { ChartsPluginStart } from '@kbn/charts-plugin/public';
 import {
   AppMountParameters,
   AppNavLinkStatus,
@@ -19,58 +20,58 @@ import {
   PluginInitializerContext,
 } from '@kbn/core/public';
 import type {
-  DataPublicPluginStart,
   DataPublicPluginSetup,
+  DataPublicPluginStart,
 } from '@kbn/data-plugin/public';
-import { LensPublicStart } from '@kbn/lens-plugin/public';
-import type { UnifiedSearchPublicPluginStart } from '@kbn/unified-search-plugin/public';
-import type { ExploratoryViewPublicSetup } from '@kbn/exploratory-view-plugin/public';
+import { DataViewsPublicPluginStart } from '@kbn/data-views-plugin/public';
+import {
+  DiscoverSetup,
+  DiscoverStart,
+} from '@kbn/discover-plugin/public/plugin';
 import type { EmbeddableStart } from '@kbn/embeddable-plugin/public';
-import type { HomePublicPluginSetup } from '@kbn/home-plugin/public';
-import { Start as InspectorPluginStart } from '@kbn/inspector-plugin/public';
-import type {
-  PluginSetupContract as AlertingPluginPublicSetup,
-  PluginStartContract as AlertingPluginPublicStart,
-} from '@kbn/alerting-plugin/public';
-import type { IStorageWrapper } from '@kbn/kibana-utils-plugin/public';
+import type { ExploratoryViewPublicSetup } from '@kbn/exploratory-view-plugin/public';
 import type { FeaturesPluginSetup } from '@kbn/features-plugin/public';
+import { FieldFormatsStart } from '@kbn/field-formats-plugin/public';
 import type { FleetStart } from '@kbn/fleet-plugin/public';
+import type { HomePublicPluginSetup } from '@kbn/home-plugin/public';
+import { i18n } from '@kbn/i18n';
+import { InfraClientStartExports } from '@kbn/infra-plugin/public';
+import { Start as InspectorPluginStart } from '@kbn/inspector-plugin/public';
+import type { IStorageWrapper } from '@kbn/kibana-utils-plugin/public';
+import { LensPublicStart } from '@kbn/lens-plugin/public';
+import { LicenseManagementUIPluginSetup } from '@kbn/license-management-plugin/public';
 import type { LicensingPluginSetup } from '@kbn/licensing-plugin/public';
 import type { MapsStartApi } from '@kbn/maps-plugin/public';
 import type { MlPluginSetup, MlPluginStart } from '@kbn/ml-plugin/public';
-import type { SharePluginSetup } from '@kbn/share-plugin/public';
-import type {
-  ObservabilitySharedPluginSetup,
-  ObservabilitySharedPluginStart,
-} from '@kbn/observability-shared-plugin/public';
+import type { ObservabilityAIAssistantPluginStart } from '@kbn/observability-ai-assistant-plugin/public';
 import {
   FetchDataParams,
   ObservabilityPublicSetup,
   ObservabilityPublicStart,
 } from '@kbn/observability-plugin/public';
-import { METRIC_TYPE } from '@kbn/observability-shared-plugin/public';
-import type {
-  TriggersAndActionsUIPublicPluginSetup,
-  TriggersAndActionsUIPublicPluginStart,
-} from '@kbn/triggers-actions-ui-plugin/public';
-import type { SecurityPluginStart } from '@kbn/security-plugin/public';
-import { SpacesPluginStart } from '@kbn/spaces-plugin/public';
-import { InfraClientStartExports } from '@kbn/infra-plugin/public';
-import { DataViewsPublicPluginStart } from '@kbn/data-views-plugin/public';
-import { ChartsPluginStart } from '@kbn/charts-plugin/public';
-import { FieldFormatsStart } from '@kbn/field-formats-plugin/public';
-import { UiActionsStart, UiActionsSetup } from '@kbn/ui-actions-plugin/public';
 import { ObservabilityTriggerId } from '@kbn/observability-shared-plugin/common';
-import { LicenseManagementUIPluginSetup } from '@kbn/license-management-plugin/public';
+import type {
+  ObservabilitySharedPluginSetup,
+  ObservabilitySharedPluginStart,
+} from '@kbn/observability-shared-plugin/public';
+import { METRIC_TYPE } from '@kbn/observability-shared-plugin/public';
 import {
   ProfilingPluginSetup,
   ProfilingPluginStart,
 } from '@kbn/profiling-plugin/public';
-import {
-  DiscoverStart,
-  DiscoverSetup,
-} from '@kbn/discover-plugin/public/plugin';
-import type { ObservabilityAIAssistantPluginStart } from '@kbn/observability-ai-assistant-plugin/public';
+import type { SecurityPluginStart } from '@kbn/security-plugin/public';
+import type { SharePluginSetup } from '@kbn/share-plugin/public';
+import { SpacesPluginStart } from '@kbn/spaces-plugin/public';
+import type {
+  TriggersAndActionsUIPublicPluginSetup,
+  TriggersAndActionsUIPublicPluginStart,
+} from '@kbn/triggers-actions-ui-plugin/public';
+import { UiActionsSetup, UiActionsStart } from '@kbn/ui-actions-plugin/public';
+import type { UnifiedSearchPublicPluginStart } from '@kbn/unified-search-plugin/public';
+import { UsageCollectionStart } from '@kbn/usage-collection-plugin/public';
+import { from } from 'rxjs';
+import { map } from 'rxjs/operators';
+import type { ConfigSchema } from '.';
 import { registerApmRuleTypes } from './components/alerting/rule_types/register_apm_rule_types';
 import {
   getApmEnrollmentFlyoutData,
@@ -80,7 +81,6 @@ import { getLazyApmAgentsTabExtension } from './components/fleet_integration/laz
 import { getLazyAPMPolicyCreateExtension } from './components/fleet_integration/lazy_apm_policy_create_extension';
 import { getLazyAPMPolicyEditExtension } from './components/fleet_integration/lazy_apm_policy_edit_extension';
 import { featureCatalogueEntry } from './feature_catalogue_entry';
-import type { ConfigSchema } from '.';
 import { APMServiceDetailLocator } from './locator/service_detail_locator';
 
 export type ApmPluginSetup = ReturnType<ApmPlugin['setup']>;
@@ -413,6 +413,20 @@ export class ApmPlugin implements Plugin<ApmPluginSetup, ApmPluginStart> {
 
   public start(core: CoreStart, plugins: ApmPluginStartDeps) {
     const { fleet } = plugins;
+
+    plugins.observabilityAIAssistant.register(
+      async ({ signal, registerFunction, registerContext }) => {
+        const mod = await import('./components/assistant_functions');
+
+        mod.registerAssistantFunctions({
+          coreStart: core,
+          pluginsStart: plugins,
+          registerFunction,
+          registerContext,
+        });
+      }
+    );
+
     if (fleet) {
       const agentEnrollmentExtensionData = getApmEnrollmentFlyoutData();
 

x-pack/plugins/apm/server/lib/connections/get_connection_stats/get_destination_map.ts

@@ -208,6 +208,7 @@ export const getDestinationMap = ({
           environment: mergedDestination.environment,
           id: objectHash({ serviceName: mergedDestination.serviceName }),
           type: NodeType.service,
+          dependencyName: mergedDestination.dependencyName,
         };
       } else {
         node = {

x-pack/plugins/apm/server/routes/apm_routes/get_global_apm_server_route_repository.ts

@@ -44,6 +44,7 @@ import { suggestionsRouteRepository } from '../suggestions/route';
 import { timeRangeMetadataRoute } from '../time_range_metadata/route';
 import { traceRouteRepository } from '../traces/route';
 import { transactionRouteRepository } from '../transactions/route';
+import { assistantRouteRepository } from '../assistant_functions/route';
 
 function getTypedGlobalApmServerRouteRepository() {
   const repository = {
@@ -81,6 +82,7 @@ function getTypedGlobalApmServerRouteRepository() {
     ...agentExplorerRouteRepository,
     ...mobileRouteRepository,
     ...diagnosticsRepository,
+    ...assistantRouteRepository,
   };
 
   return repository;

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_correlation_values/index.ts

@@ -0,0 +1,176 @@
+/*
+ * 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 { AggregationsSignificantTermsAggregation } from '@elastic/elasticsearch/lib/api/types';
+import { ProcessorEvent } from '@kbn/observability-plugin/common';
+import { kqlQuery, rangeQuery } from '@kbn/observability-plugin/server';
+import * as t from 'io-ts';
+import { CorrelationsEventType } from '../../../../common/assistant/constants';
+import {
+  SERVICE_NAME,
+  SERVICE_NODE_NAME,
+  SPAN_DESTINATION_SERVICE_RESOURCE,
+  SPAN_NAME,
+  TRANSACTION_NAME,
+  TRANSACTION_RESULT,
+} from '../../../../common/es_fields/apm';
+import { termQuery } from '../../../../common/utils/term_query';
+import {
+  APMEventClient,
+  APMEventESSearchRequest,
+} from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { environmentRt } from '../../default_api_types';
+
+const setRt = t.intersection([
+  t.type({
+    start: t.string,
+    end: t.string,
+    'service.name': t.string,
+    label: t.string,
+  }),
+  t.partial({
+    filter: t.string,
+    'service.environment': environmentRt.props.environment,
+  }),
+]);
+
+export const correlationValuesRouteRt = t.type({
+  sets: t.array(
+    t.type({
+      foreground: setRt,
+      background: setRt,
+      event: t.union([
+        t.literal(CorrelationsEventType.Transaction),
+        t.literal(CorrelationsEventType.ExitSpan),
+        t.literal(CorrelationsEventType.Error),
+      ]),
+    })
+  ),
+});
+
+export interface CorrelationValue {
+  foreground: string;
+  background: string;
+  fieldName: string;
+  fields: Array<{ value: string; score: number }>;
+}
+
+export async function getApmCorrelationValues({
+  arguments: args,
+  apmEventClient,
+}: {
+  arguments: t.TypeOf<typeof correlationValuesRouteRt>;
+  apmEventClient: APMEventClient;
+}): Promise<CorrelationValue[]> {
+  const getQueryForSet = (set: t.TypeOf<typeof setRt>) => {
+    const start = datemath.parse(set.start)?.valueOf()!;
+    const end = datemath.parse(set.end)?.valueOf()!;
+
+    return {
+      bool: {
+        filter: [
+          ...rangeQuery(start, end),
+          ...termQuery(SERVICE_NAME, set['service.name']),
+          ...kqlQuery(set.filter),
+        ],
+      },
+    };
+  };
+
+  const allCorrelations = await Promise.all(
+    args.sets.map(async (set) => {
+      const query = getQueryForSet(set.foreground);
+
+      let apm: APMEventESSearchRequest['apm'];
+
+      let fields: string[] = [];
+
+      switch (set.event) {
+        case CorrelationsEventType.Transaction:
+          apm = {
+            events: [ProcessorEvent.transaction],
+          };
+          fields = [TRANSACTION_NAME, SERVICE_NODE_NAME, TRANSACTION_RESULT];
+          break;
+
+        case CorrelationsEventType.ExitSpan:
+          apm = {
+            events: [ProcessorEvent.span],
+          };
+          fields = [SPAN_NAME, SPAN_DESTINATION_SERVICE_RESOURCE];
+          query.bool.filter.push({
+            exists: {
+              field: SPAN_DESTINATION_SERVICE_RESOURCE,
+            },
+          });
+          break;
+
+        case CorrelationsEventType.Error:
+          apm = {
+            events: [ProcessorEvent.error],
+          };
+          fields = ['error.grouping_name'];
+          break;
+      }
+
+      const sigTermsAggs: Record<
+        string,
+        { significant_terms: AggregationsSignificantTermsAggregation }
+      > = {};
+
+      fields.forEach((field) => {
+        sigTermsAggs[field] = {
+          significant_terms: {
+            field,
+            background_filter: getQueryForSet(set.background),
+            gnd: {
+              background_is_superset: false,
+            },
+          },
+        };
+      });
+
+      const response = await apmEventClient.search('get_significant_terms', {
+        apm,
+        body: {
+          size: 0,
+          track_total_hits: false,
+          query,
+          aggs: sigTermsAggs,
+        },
+      });
+
+      const correlations: Array<{
+        foreground: string;
+        background: string;
+        fieldName: string;
+        fields: Array<{ value: string; score: number }>;
+      }> = [];
+
+      if (!response.aggregations) {
+        return { correlations: [] };
+      }
+
+      // eslint-disable-next-line guard-for-in
+      for (const fieldName in response.aggregations) {
+        correlations.push({
+          foreground: set.foreground.label,
+          background: set.background.label,
+          fieldName,
+          fields: response.aggregations[fieldName].buckets.map((bucket) => ({
+            score: bucket.score,
+            value: String(bucket.key),
+          })),
+        });
+      }
+
+      return { correlations };
+    })
+  );
+
+  return allCorrelations.flatMap((_) => _.correlations);
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_downstream_dependencies/index.ts

@@ -0,0 +1,79 @@
+/*
+ * 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 * as t from 'io-ts';
+import { termQuery } from '@kbn/observability-plugin/server';
+import { ENVIRONMENT_ALL } from '../../../../common/environment_filter_values';
+import { SERVICE_NAME } from '../../../../common/es_fields/apm';
+import { environmentQuery } from '../../../../common/utils/environment_query';
+import { getDestinationMap } from '../../../lib/connections/get_connection_stats/get_destination_map';
+import { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { NodeType } from '../../../../common/connections';
+
+export const downstreamDependenciesRouteRt = t.intersection([
+  t.type({
+    'service.name': t.string,
+    start: t.string,
+    end: t.string,
+  }),
+  t.partial({
+    'service.environment': t.string,
+  }),
+]);
+
+export interface APMDownstreamDependency {
+  'service.name'?: string | undefined;
+  'span.destination.service.resource': string;
+  'span.type'?: string | undefined;
+  'span.subtype'?: string | undefined;
+}
+
+export async function getAssistantDownstreamDependencies({
+  arguments: args,
+  apmEventClient,
+}: {
+  arguments: t.TypeOf<typeof downstreamDependenciesRouteRt>;
+  apmEventClient: APMEventClient;
+}): Promise<APMDownstreamDependency[]> {
+  const start = datemath.parse(args.start)?.valueOf()!;
+  const end = datemath.parse(args.end)?.valueOf()!;
+
+  const map = await getDestinationMap({
+    start,
+    end,
+    apmEventClient,
+    filter: [
+      ...termQuery(SERVICE_NAME, args['service.name']),
+      ...environmentQuery(args['service.environment'] ?? ENVIRONMENT_ALL.value),
+    ],
+  });
+
+  const items: Array<{
+    'service.name'?: string;
+    'span.destination.service.resource': string;
+    'span.type'?: string;
+    'span.subtype'?: string;
+  }> = [];
+
+  for (const [_, node] of map) {
+    if (node.type === NodeType.service) {
+      items.push({
+        'service.name': node.serviceName,
+        // this should be set, as it's a downstream dependency, and there should be a connection
+        'span.destination.service.resource': node.dependencyName!,
+      });
+    } else {
+      items.push({
+        'span.destination.service.resource': node.dependencyName,
+        'span.type': node.spanType,
+        'span.subtype': node.spanSubtype,
+      });
+    }
+  }
+
+  return items;
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_error_document/index.ts

@@ -0,0 +1,74 @@
+/*
+ * 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 * as t from 'io-ts';
+import { rangeQuery } from '@kbn/observability-plugin/server';
+import datemath from '@elastic/datemath';
+import { pick } from 'lodash';
+import { ApmDocumentType } from '../../../../common/document_type';
+import { RollupInterval } from '../../../../common/rollup';
+import { termQuery } from '../../../../common/utils/term_query';
+import { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { APMError } from '../../../../typings/es_schemas/ui/apm_error';
+
+export const errorRouteRt = t.type({
+  start: t.string,
+  end: t.string,
+  'error.grouping_name': t.string,
+});
+
+export async function getApmErrorDocument({
+  arguments: args,
+  apmEventClient,
+}: {
+  arguments: t.TypeOf<typeof errorRouteRt>;
+  apmEventClient: APMEventClient;
+}) {
+  const start = datemath.parse(args.start)?.valueOf()!;
+  const end = datemath.parse(args.end)?.valueOf()!;
+
+  const response = await apmEventClient.search('get_error', {
+    apm: {
+      sources: [
+        {
+          documentType: ApmDocumentType.ErrorEvent,
+          rollupInterval: RollupInterval.None,
+        },
+      ],
+    },
+    body: {
+      track_total_hits: false,
+      size: 1,
+      terminate_after: 1,
+      query: {
+        bool: {
+          filter: [
+            ...rangeQuery(start, end),
+            ...termQuery('error.grouping_name', args['error.grouping_name']),
+          ],
+        },
+      },
+    },
+  });
+
+  const error = response.hits.hits[0]?._source as APMError;
+
+  if (!error) {
+    return undefined;
+  }
+
+  return pick(
+    error,
+    'message',
+    'error',
+    '@timestamp',
+    'transaction.name',
+    'transaction.type',
+    'span.name',
+    'span.type',
+    'span.subtype'
+  );
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_service_summary/index.ts

@@ -0,0 +1,318 @@
+/*
+ * 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 { ElasticsearchClient, Logger } from '@kbn/core/server';
+import {
+  rangeQuery,
+  ScopedAnnotationsClient,
+} from '@kbn/observability-plugin/server';
+import {
+  ALERT_RULE_PRODUCER,
+  ALERT_STATUS,
+  ALERT_STATUS_ACTIVE,
+} from '@kbn/rule-registry-plugin/common/technical_rule_data_field_names';
+import * as t from 'io-ts';
+import { compact, keyBy } from 'lodash';
+import {
+  ApmMlDetectorType,
+  getApmMlDetectorType,
+} from '../../../../common/anomaly_detection/apm_ml_detectors';
+import { ENVIRONMENT_ALL } from '../../../../common/environment_filter_values';
+import { Environment } from '../../../../common/environment_rt';
+import { SERVICE_NAME } from '../../../../common/es_fields/apm';
+import { asMutableArray } from '../../../../common/utils/as_mutable_array';
+import { environmentQuery } from '../../../../common/utils/environment_query';
+import { maybe } from '../../../../common/utils/maybe';
+import { termQuery } from '../../../../common/utils/term_query';
+import { anomalySearch } from '../../../lib/anomaly_detection/anomaly_search';
+import { apmMlAnomalyQuery } from '../../../lib/anomaly_detection/apm_ml_anomaly_query';
+import { apmMlJobsQuery } from '../../../lib/anomaly_detection/apm_ml_jobs_query';
+import { getMlJobsWithAPMGroup } from '../../../lib/anomaly_detection/get_ml_jobs_with_apm_group';
+import { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { ApmAlertsClient } from '../../../lib/helpers/get_apm_alerts_client';
+import { MlClient } from '../../../lib/helpers/get_ml_client';
+import { getServiceAnnotations } from '../../services/annotations';
+import { getServiceMetadataDetails } from '../../services/get_service_metadata_details';
+
+export const serviceSummaryRouteRt = t.intersection([
+  t.type({
+    'service.name': t.string,
+    start: t.string,
+    end: t.string,
+  }),
+  t.partial({
+    'service.environment': t.string,
+    'transaction.type': t.string,
+  }),
+]);
+
+async function getAnomalies({
+  serviceName,
+  transactionType,
+  environment,
+  start,
+  end,
+  mlClient,
+  logger,
+}: {
+  serviceName: string;
+  transactionType?: string;
+  environment?: string;
+  start: number;
+  end: number;
+  mlClient?: MlClient;
+  logger: Logger;
+}) {
+  if (!mlClient) {
+    return [];
+  }
+
+  const mlJobs = (
+    await getMlJobsWithAPMGroup(mlClient.anomalyDetectors)
+  ).filter((job) => job.environment !== environment);
+
+  if (!mlJobs.length) {
+    return [];
+  }
+
+  const anomaliesResponse = await anomalySearch(
+    mlClient.mlSystem.mlAnomalySearch,
+    {
+      body: {
+        size: 0,
+        query: {
+          bool: {
+            filter: [
+              ...apmMlAnomalyQuery({
+                serviceName,
+                transactionType,
+              }),
+              ...rangeQuery(start, end, 'timestamp'),
+              ...apmMlJobsQuery(mlJobs),
+            ],
+          },
+        },
+        aggs: {
+          by_timeseries_id: {
+            composite: {
+              size: 5000,
+              sources: asMutableArray([
+                {
+                  jobId: {
+                    terms: {
+                      field: 'job_id',
+                    },
+                  },
+                },
+                {
+                  detectorIndex: {
+                    terms: {
+                      field: 'detector_index',
+                    },
+                  },
+                },
+                {
+                  serviceName: {
+                    terms: {
+                      field: 'partition_field_value',
+                    },
+                  },
+                },
+                {
+                  transactionType: {
+                    terms: {
+                      field: 'by_field_value',
+                    },
+                  },
+                },
+              ] as const),
+            },
+            aggs: {
+              record_scores: {
+                filter: {
+                  term: {
+                    result_type: 'record',
+                  },
+                },
+                aggs: {
+                  top_anomaly: {
+                    top_metrics: {
+                      metrics: asMutableArray([
+                        { field: 'record_score' },
+                        { field: 'actual' },
+                        { field: 'timestamp' },
+                      ] as const),
+                      size: 1,
+                      sort: {
+                        record_score: 'desc',
+                      },
+                    },
+                  },
+                },
+              },
+              model_lower: {
+                min: {
+                  field: 'model_lower',
+                },
+              },
+              model_upper: {
+                max: {
+                  field: 'model_upper',
+                },
+              },
+            },
+          },
+        },
+      },
+    }
+  );
+
+  const jobsById = keyBy(mlJobs, (job) => job.jobId);
+
+  const anomalies =
+    anomaliesResponse.aggregations?.by_timeseries_id.buckets.map((bucket) => {
+      const jobId = bucket.key.jobId as string;
+      const job = maybe(jobsById[jobId]);
+
+      if (!job) {
+        logger.warn(`Could not find job for id ${jobId}`);
+        return undefined;
+      }
+
+      const type = getApmMlDetectorType(Number(bucket.key.detectorIndex));
+
+      // ml failure rate is stored as 0-100, we calculate failure rate as 0-1
+      const divider = type === ApmMlDetectorType.txFailureRate ? 100 : 1;
+
+      const metrics = bucket.record_scores.top_anomaly.top[0]?.metrics;
+
+      if (!metrics) {
+        return undefined;
+      }
+
+      return {
+        '@timestamp': new Date(metrics.timestamp as number).toISOString(),
+        metricName: type.replace('tx', 'transaction'),
+        'service.name': bucket.key.serviceName as string,
+        'service.environment': job.environment,
+        'transaction.type': bucket.key.transactionType as string,
+        anomalyScore: metrics.record_score,
+        actualValue: Number(metrics.actual) / divider,
+        expectedBoundsLower: Number(bucket.model_lower.value) / divider,
+        expectedBoundsUpper: Number(bucket.model_upper.value) / divider,
+      };
+    });
+
+  return compact(anomalies);
+}
+
+export interface ServiceSummary {
+  'service.name': string;
+  'agent.name'?: string;
+  'service.version'?: string[];
+  'language.name'?: string;
+  'service.framework'?: string;
+  instances: number;
+  anomalies: Array<{
+    '@timestamp': string;
+    metricName: string;
+    'service.name': string;
+    'service.environment': Environment;
+    'transaction.type': string;
+    anomalyScore: string | number | null;
+    actualValue: number;
+    expectedBoundsLower: number;
+    expectedBoundsUpper: number;
+  }>;
+  alerts: Array<{ type?: string; started: string }>;
+  deployments: Array<{ '@timestamp': string }>;
+}
+
+export async function getApmServiceSummary({
+  arguments: args,
+  apmEventClient,
+  mlClient,
+  esClient,
+  annotationsClient,
+  apmAlertsClient,
+  logger,
+}: {
+  arguments: t.TypeOf<typeof serviceSummaryRouteRt>;
+  apmEventClient: APMEventClient;
+  mlClient?: MlClient;
+  esClient: ElasticsearchClient;
+  annotationsClient?: ScopedAnnotationsClient;
+  apmAlertsClient: ApmAlertsClient;
+  logger: Logger;
+}): Promise<ServiceSummary> {
+  const start = datemath.parse(args.start)?.valueOf()!;
+  const end = datemath.parse(args.end)?.valueOf()!;
+
+  const serviceName = args['service.name'];
+  const environment = args['service.environment'] || ENVIRONMENT_ALL.value;
+  const transactionType = args['transaction.type'];
+
+  const [metadataDetails, anomalies, annotations, alerts] = await Promise.all([
+    getServiceMetadataDetails({
+      apmEventClient,
+      start,
+      end,
+      serviceName,
+    }),
+    getAnomalies({
+      serviceName,
+      start,
+      end,
+      environment,
+      mlClient,
+      logger,
+      transactionType,
+    }),
+    getServiceAnnotations({
+      apmEventClient,
+      start,
+      end,
+      searchAggregatedTransactions: true,
+      client: esClient,
+      annotationsClient,
+      environment,
+      logger,
+      serviceName,
+    }),
+    apmAlertsClient.search({
+      size: 100,
+      query: {
+        bool: {
+          filter: [
+            ...termQuery(ALERT_RULE_PRODUCER, 'apm'),
+            ...termQuery(ALERT_STATUS, ALERT_STATUS_ACTIVE),
+            ...rangeQuery(start, end),
+            ...termQuery(SERVICE_NAME, serviceName),
+            ...environmentQuery(environment),
+          ],
+        },
+      },
+    }),
+  ]);
+
+  return {
+    'service.name': serviceName,
+    'agent.name': metadataDetails.service?.agent.name,
+    'service.version': metadataDetails.service?.versions,
+    'language.name': metadataDetails.service?.agent.name,
+    'service.framework': metadataDetails.service?.framework,
+    instances: metadataDetails.container?.totalNumberInstances ?? 1,
+    anomalies,
+    alerts: alerts.hits.hits.map((alert) => ({
+      type: alert._source?.['kibana.alert.rule.type'],
+      started: new Date(alert._source?.['kibana.alert.start']!).toISOString(),
+    })),
+    deployments: annotations.annotations.map((annotation) => ({
+      '@timestamp': new Date(annotation['@timestamp']).toISOString(),
+    })),
+  };
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/fetch_timeseries.ts

@@ -0,0 +1,136 @@
+/*
+ * 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 {
+  AggregationsAggregationContainer,
+  QueryDslQueryContainer,
+} from '@elastic/elasticsearch/lib/api/types';
+import { AggregationResultOf, AggregationResultOfMap } from '@kbn/es-types';
+import { Unionize } from 'utility-types';
+import { ApmDocumentType } from '../../../../common/document_type';
+import { RollupInterval } from '../../../../common/rollup';
+import { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+
+type ChangePointResult = AggregationResultOf<{ change_point: any }, unknown>;
+
+type ValueAggregationMap = Record<
+  'value',
+  Unionize<
+    Pick<
+      Required<AggregationsAggregationContainer>,
+      'min' | 'max' | 'sum' | 'bucket_script' | 'avg'
+    >
+  >
+>;
+
+interface ApmFetchedTimeseries<T extends ValueAggregationMap> {
+  groupBy: string;
+  data: Array<
+    {
+      key: number;
+      key_as_string: string;
+      doc_count: number;
+    } & AggregationResultOfMap<T, unknown>
+  >;
+  change_point: ChangePointResult;
+  value: number | null;
+  unit: string;
+}
+
+export interface FetchSeriesProps<T extends ValueAggregationMap> {
+  apmEventClient: APMEventClient;
+  operationName: string;
+  documentType: ApmDocumentType;
+  rollupInterval: RollupInterval;
+  intervalString: string;
+  start: number;
+  end: number;
+  filter?: QueryDslQueryContainer[];
+  groupBy?: string;
+  aggs: T;
+  unit: 'ms' | 'rpm' | '%';
+}
+
+export async function fetchSeries<T extends ValueAggregationMap>({
+  apmEventClient,
+  operationName,
+  documentType,
+  rollupInterval,
+  intervalString,
+  start,
+  end,
+  filter,
+  groupBy,
+  aggs,
+  unit,
+}: FetchSeriesProps<T>): Promise<Array<ApmFetchedTimeseries<T>>> {
+  const response = await apmEventClient.search(operationName, {
+    apm: {
+      sources: [{ documentType, rollupInterval }],
+    },
+    body: {
+      size: 0,
+      track_total_hits: false,
+      query: {
+        bool: {
+          filter,
+        },
+      },
+      aggs: {
+        groupBy: {
+          ...(groupBy
+            ? {
+                terms: {
+                  field: groupBy,
+                  size: 20,
+                },
+              }
+            : {
+                terms: {
+                  field: 'non_existing_field',
+                  missing: '',
+                },
+              }),
+          aggs: {
+            ...aggs,
+            timeseries: {
+              date_histogram: {
+                field: '@timestamp',
+                fixed_interval: intervalString,
+                min_doc_count: 0,
+                extended_bounds: { min: start, max: end },
+              },
+              aggs,
+            },
+            change_point: {
+              change_point: {
+                buckets_path: 'timeseries>value',
+              },
+            },
+          },
+        },
+      },
+    },
+  });
+
+  if (!response.aggregations?.groupBy) {
+    return [];
+  }
+
+  return response.aggregations.groupBy.buckets.map((bucket) => {
+    return {
+      groupBy: bucket.key_as_string || String(bucket.key),
+      data: bucket.timeseries.buckets,
+      value:
+        bucket.value?.value === undefined || bucket.value?.value === null
+          ? null
+          : Math.round(bucket.value.value),
+      change_point: bucket.change_point,
+      unit,
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/get_error_event_rate.ts

@@ -0,0 +1,76 @@
+/*
+ * 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 { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { rangeQuery } from '@kbn/observability-plugin/server';
+import { ApmDocumentType } from '../../../../common/document_type';
+import { RollupInterval } from '../../../../common/rollup';
+import type { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { fetchSeries } from './fetch_timeseries';
+
+export async function getErrorEventRate({
+  apmEventClient,
+  start,
+  end,
+  intervalString,
+  bucketSize,
+  filter,
+}: {
+  apmEventClient: APMEventClient;
+  start: number;
+  end: number;
+  intervalString: string;
+  bucketSize: number;
+  filter: QueryDslQueryContainer[];
+}) {
+  const bucketSizeInMinutes = bucketSize / 60;
+  const rangeInMinutes = (end - start) / 1000 / 60;
+
+  return (
+    await fetchSeries({
+      apmEventClient,
+      start,
+      end,
+      operationName: 'assistant_get_error_event_rate',
+      unit: 'rpm',
+      documentType: ApmDocumentType.ErrorEvent,
+      rollupInterval: RollupInterval.None,
+      intervalString,
+      filter: filter.concat(...rangeQuery(start, end)),
+      aggs: {
+        value: {
+          bucket_script: {
+            buckets_path: {
+              count: '_count',
+            },
+            script: {
+              lang: 'painless',
+              params: {
+                bucketSizeInMinutes,
+              },
+              source: 'params.count / params.bucketSizeInMinutes',
+            },
+          },
+        },
+      },
+    })
+  ).map((fetchedSerie) => {
+    return {
+      ...fetchedSerie,
+      value:
+        fetchedSerie.value !== null
+          ? fetchedSerie.value / rangeInMinutes
+          : null,
+      data: fetchedSerie.data.map((bucket) => {
+        return {
+          x: bucket.key,
+          y: bucket.value?.value as number,
+        };
+      }),
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/get_exit_span_failure_rate.ts

@@ -0,0 +1,107 @@
+/*
+ * 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 { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { rangeQuery, termQuery } from '@kbn/observability-plugin/server';
+import { ApmDocumentType } from '../../../../common/document_type';
+import {
+  EVENT_OUTCOME,
+  SPAN_DESTINATION_SERVICE_RESOURCE,
+  SPAN_DESTINATION_SERVICE_RESPONSE_TIME_COUNT,
+} from '../../../../common/es_fields/apm';
+import { EventOutcome } from '../../../../common/event_outcome';
+import { RollupInterval } from '../../../../common/rollup';
+import type { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { fetchSeries } from './fetch_timeseries';
+
+export async function getExitSpanFailureRate({
+  apmEventClient,
+  start,
+  end,
+  intervalString,
+  filter,
+  spanDestinationServiceResource,
+}: {
+  apmEventClient: APMEventClient;
+  start: number;
+  end: number;
+  intervalString: string;
+  bucketSize: number;
+  filter: QueryDslQueryContainer[];
+  spanDestinationServiceResource?: string;
+}) {
+  return (
+    await fetchSeries({
+      apmEventClient,
+      start,
+      end,
+      operationName: 'assistant_get_exit_span_failure_rate',
+      unit: '%',
+      documentType: ApmDocumentType.ServiceDestinationMetric,
+      rollupInterval: RollupInterval.OneMinute,
+      intervalString,
+      filter: filter.concat(
+        ...rangeQuery(start, end),
+        ...termQuery(
+          SPAN_DESTINATION_SERVICE_RESOURCE,
+          spanDestinationServiceResource
+        )
+      ),
+      groupBy: SPAN_DESTINATION_SERVICE_RESOURCE,
+      aggs: {
+        successful: {
+          filter: {
+            terms: {
+              [EVENT_OUTCOME]: [EventOutcome.success],
+            },
+          },
+          aggs: {
+            count: {
+              sum: {
+                field: SPAN_DESTINATION_SERVICE_RESPONSE_TIME_COUNT,
+              },
+            },
+          },
+        },
+        successful_or_failed: {
+          filter: {
+            terms: {
+              [EVENT_OUTCOME]: [EventOutcome.success, EventOutcome.failure],
+            },
+          },
+          aggs: {
+            count: {
+              sum: {
+                field: SPAN_DESTINATION_SERVICE_RESPONSE_TIME_COUNT,
+              },
+            },
+          },
+        },
+        value: {
+          bucket_script: {
+            buckets_path: {
+              successful_or_failed: `successful_or_failed>count`,
+              successful: `successful>count`,
+            },
+            script:
+              '100 * (1 - (params.successful / params.successful_or_failed))',
+          },
+        },
+      },
+    })
+  ).map((fetchedSerie) => {
+    return {
+      ...fetchedSerie,
+      data: fetchedSerie.data.map((bucket) => {
+        return {
+          x: bucket.key,
+          y: bucket.value?.value as number | null,
+        };
+      }),
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/get_exit_span_latency.ts

@@ -0,0 +1,87 @@
+/*
+ * 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 { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { rangeQuery, termQuery } from '@kbn/observability-plugin/server';
+import { ApmDocumentType } from '../../../../common/document_type';
+import {
+  SPAN_DESTINATION_SERVICE_RESOURCE,
+  SPAN_DESTINATION_SERVICE_RESPONSE_TIME_COUNT,
+  SPAN_DESTINATION_SERVICE_RESPONSE_TIME_SUM,
+} from '../../../../common/es_fields/apm';
+import { RollupInterval } from '../../../../common/rollup';
+import { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { fetchSeries } from './fetch_timeseries';
+
+export async function getExitSpanLatency({
+  apmEventClient,
+  start,
+  end,
+  intervalString,
+  filter,
+  spanDestinationServiceResource,
+}: {
+  apmEventClient: APMEventClient;
+  start: number;
+  end: number;
+  intervalString: string;
+  bucketSize: number;
+  filter: QueryDslQueryContainer[];
+  spanDestinationServiceResource?: string;
+}) {
+  return (
+    await fetchSeries({
+      apmEventClient,
+      start,
+      end,
+      operationName: 'assistant_get_exit_span_latency',
+      unit: 'rpm',
+      documentType: ApmDocumentType.ServiceDestinationMetric,
+      rollupInterval: RollupInterval.OneMinute,
+      intervalString,
+      filter: filter.concat(
+        ...rangeQuery(start, end),
+        ...termQuery(
+          SPAN_DESTINATION_SERVICE_RESOURCE,
+          spanDestinationServiceResource
+        )
+      ),
+      groupBy: SPAN_DESTINATION_SERVICE_RESOURCE,
+      aggs: {
+        count: {
+          sum: {
+            field: SPAN_DESTINATION_SERVICE_RESPONSE_TIME_COUNT,
+          },
+        },
+        latency: {
+          sum: {
+            field: SPAN_DESTINATION_SERVICE_RESPONSE_TIME_SUM,
+          },
+        },
+        value: {
+          bucket_script: {
+            buckets_path: {
+              latency: 'latency',
+              count: 'count',
+            },
+            script: '(params.latency / params.count) / 1000',
+          },
+        },
+      },
+    })
+  ).map((fetchedSerie) => {
+    return {
+      ...fetchedSerie,
+      data: fetchedSerie.data.map((bucket) => {
+        return {
+          x: bucket.key,
+          y: bucket.value?.value as number | null,
+        };
+      }),
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/get_exit_span_throughput.ts

@@ -0,0 +1,86 @@
+/*
+ * 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 { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { rangeQuery, termQuery } from '@kbn/observability-plugin/server';
+import { ApmDocumentType } from '../../../../common/document_type';
+import { SPAN_DESTINATION_SERVICE_RESOURCE } from '../../../../common/es_fields/apm';
+import { RollupInterval } from '../../../../common/rollup';
+import type { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { fetchSeries } from './fetch_timeseries';
+
+export async function getExitSpanThroughput({
+  apmEventClient,
+  start,
+  end,
+  intervalString,
+  bucketSize,
+  filter,
+  spanDestinationServiceResource,
+}: {
+  apmEventClient: APMEventClient;
+  start: number;
+  end: number;
+  intervalString: string;
+  bucketSize: number;
+  filter: QueryDslQueryContainer[];
+  spanDestinationServiceResource?: string;
+}) {
+  const bucketSizeInMinutes = bucketSize / 60;
+  const rangeInMinutes = (end - start) / 1000 / 60;
+
+  return (
+    await fetchSeries({
+      apmEventClient,
+      start,
+      end,
+      operationName: 'assistant_get_exit_span_throughput',
+      unit: 'rpm',
+      documentType: ApmDocumentType.ServiceDestinationMetric,
+      rollupInterval: RollupInterval.OneMinute,
+      intervalString,
+      filter: filter.concat(
+        ...rangeQuery(start, end),
+        ...termQuery(
+          SPAN_DESTINATION_SERVICE_RESOURCE,
+          spanDestinationServiceResource
+        )
+      ),
+      groupBy: SPAN_DESTINATION_SERVICE_RESOURCE,
+      aggs: {
+        value: {
+          bucket_script: {
+            buckets_path: {
+              count: '_count',
+            },
+            script: {
+              lang: 'painless',
+              params: {
+                bucketSizeInMinutes,
+              },
+              source: 'params.count / params.bucketSizeInMinutes',
+            },
+          },
+        },
+      },
+    })
+  ).map((fetchedSerie) => {
+    return {
+      ...fetchedSerie,
+      value:
+        fetchedSerie.value !== null
+          ? fetchedSerie.value / rangeInMinutes
+          : null,
+      data: fetchedSerie.data.map((bucket) => {
+        return {
+          x: bucket.key,
+          y: bucket.value?.value as number,
+        };
+      }),
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/get_transaction_failure_rate.ts

@@ -0,0 +1,73 @@
+/*
+ * 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 { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { rangeQuery, termQuery } from '@kbn/observability-plugin/server';
+import { ApmDocumentType } from '../../../../common/document_type';
+import { TRANSACTION_TYPE } from '../../../../common/es_fields/apm';
+import { RollupInterval } from '../../../../common/rollup';
+import type { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { getOutcomeAggregation } from '../../../lib/helpers/transaction_error_rate';
+import { fetchSeries } from './fetch_timeseries';
+
+export async function getTransactionFailureRate({
+  apmEventClient,
+  start,
+  end,
+  intervalString,
+  filter,
+  transactionType,
+}: {
+  apmEventClient: APMEventClient;
+  start: number;
+  end: number;
+  intervalString: string;
+  bucketSize: number;
+  filter: QueryDslQueryContainer[];
+  transactionType?: string;
+}) {
+  return (
+    await fetchSeries({
+      apmEventClient,
+      start,
+      end,
+      operationName: 'assistant_get_transaction_failure_rate',
+      unit: '%',
+      documentType: ApmDocumentType.TransactionMetric,
+      rollupInterval: RollupInterval.OneMinute,
+      intervalString,
+      filter: filter.concat(
+        ...rangeQuery(start, end),
+        ...termQuery(TRANSACTION_TYPE, transactionType)
+      ),
+      groupBy: 'transaction.type',
+      aggs: {
+        ...getOutcomeAggregation(ApmDocumentType.TransactionMetric),
+        value: {
+          bucket_script: {
+            buckets_path: {
+              successful_or_failed: 'successful_or_failed>_count',
+              successful: 'successful>_count',
+            },
+            script:
+              '100 * (1 - (params.successful / params.successful_or_failed))',
+          },
+        },
+      },
+    })
+  ).map((fetchedSerie) => {
+    return {
+      ...fetchedSerie,
+      data: fetchedSerie.data.map((bucket) => {
+        return {
+          x: bucket.key,
+          y: bucket.value?.value as number | null,
+        };
+      }),
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/get_transaction_latency.ts

@@ -0,0 +1,80 @@
+/*
+ * 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 { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { rangeQuery, termQuery } from '@kbn/observability-plugin/server';
+import { ApmDocumentType } from '../../../../common/document_type';
+import {
+  TRANSACTION_DURATION_HISTOGRAM,
+  TRANSACTION_TYPE,
+} from '../../../../common/es_fields/apm';
+import { LatencyAggregationType } from '../../../../common/latency_aggregation_types';
+import { RollupInterval } from '../../../../common/rollup';
+import { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { getLatencyAggregation } from '../../../lib/helpers/latency_aggregation_type';
+import { fetchSeries } from './fetch_timeseries';
+
+export async function getTransactionLatency({
+  apmEventClient,
+  start,
+  end,
+  intervalString,
+  filter,
+  transactionType,
+  latencyAggregationType,
+}: {
+  apmEventClient: APMEventClient;
+  start: number;
+  end: number;
+  intervalString: string;
+  bucketSize: number;
+  filter: QueryDslQueryContainer[];
+  transactionType?: string;
+  latencyAggregationType: LatencyAggregationType;
+}) {
+  return (
+    await fetchSeries({
+      apmEventClient,
+      start,
+      end,
+      operationName: 'assistant_get_transaction_latencyu',
+      unit: 'rpm',
+      documentType: ApmDocumentType.TransactionMetric,
+      rollupInterval: RollupInterval.OneMinute,
+      intervalString,
+      filter: filter.concat(
+        ...rangeQuery(start, end),
+        ...termQuery(TRANSACTION_TYPE, transactionType)
+      ),
+      groupBy: 'transaction.type',
+      aggs: {
+        ...getLatencyAggregation(
+          latencyAggregationType,
+          TRANSACTION_DURATION_HISTOGRAM
+        ),
+        value: {
+          bucket_script: {
+            buckets_path: {
+              latency: 'latency',
+            },
+            script: 'params.latency / 1000',
+          },
+        },
+      },
+    })
+  ).map((fetchedSerie) => {
+    return {
+      ...fetchedSerie,
+      data: fetchedSerie.data.map((bucket) => {
+        return {
+          x: bucket.key,
+          y: bucket.value?.value as number | null,
+        };
+      }),
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/get_transaction_throughput.ts

@@ -0,0 +1,83 @@
+/*
+ * 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 { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types';
+import { rangeQuery, termQuery } from '@kbn/observability-plugin/server';
+import { ApmDocumentType } from '../../../../common/document_type';
+import { TRANSACTION_TYPE } from '../../../../common/es_fields/apm';
+import { RollupInterval } from '../../../../common/rollup';
+import type { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { fetchSeries } from './fetch_timeseries';
+
+export async function getTransactionThroughput({
+  apmEventClient,
+  start,
+  end,
+  intervalString,
+  bucketSize,
+  filter,
+  transactionType,
+}: {
+  apmEventClient: APMEventClient;
+  start: number;
+  end: number;
+  intervalString: string;
+  bucketSize: number;
+  filter: QueryDslQueryContainer[];
+  transactionType?: string;
+}) {
+  const bucketSizeInMinutes = bucketSize / 60;
+  const rangeInMinutes = (end - start) / 1000 / 60;
+
+  return (
+    await fetchSeries({
+      apmEventClient,
+      start,
+      end,
+      operationName: 'assistant_get_transaction_throughput',
+      unit: 'rpm',
+      documentType: ApmDocumentType.TransactionMetric,
+      rollupInterval: RollupInterval.OneMinute,
+      intervalString,
+      filter: filter.concat(
+        ...rangeQuery(start, end),
+        ...termQuery(TRANSACTION_TYPE, transactionType)
+      ),
+      groupBy: 'transaction.type',
+      aggs: {
+        value: {
+          bucket_script: {
+            buckets_path: {
+              count: '_count',
+            },
+            script: {
+              lang: 'painless',
+              params: {
+                bucketSizeInMinutes,
+              },
+              source: 'params.count / params.bucketSizeInMinutes',
+            },
+          },
+        },
+      },
+    })
+  ).map((fetchedSerie) => {
+    return {
+      ...fetchedSerie,
+      value:
+        fetchedSerie.value !== null
+          ? fetchedSerie.value / rangeInMinutes
+          : null,
+      data: fetchedSerie.data.map((bucket) => {
+        return {
+          x: bucket.key,
+          y: bucket.value?.value as number,
+        };
+      }),
+    };
+  });
+}

x-pack/plugins/apm/server/routes/assistant_functions/get_apm_timeseries/index.ts

@@ -0,0 +1,234 @@
+/*
+ * 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 { kqlQuery, rangeQuery } from '@kbn/observability-plugin/server';
+import * as t from 'io-ts';
+import { SERVICE_NAME } from '../../../../common/es_fields/apm';
+import { LatencyAggregationType } from '../../../../common/latency_aggregation_types';
+import { environmentQuery } from '../../../../common/utils/environment_query';
+import { getBucketSize } from '../../../../common/utils/get_bucket_size';
+import { termQuery } from '../../../../common/utils/term_query';
+import { APMEventClient } from '../../../lib/helpers/create_es_client/create_apm_event_client';
+import { environmentRt } from '../../default_api_types';
+import { getErrorEventRate } from './get_error_event_rate';
+import { getExitSpanFailureRate } from './get_exit_span_failure_rate';
+import { getExitSpanLatency } from './get_exit_span_latency';
+import { getExitSpanThroughput } from './get_exit_span_throughput';
+import { getTransactionFailureRate } from './get_transaction_failure_rate';
+import { getTransactionLatency } from './get_transaction_latency';
+import { getTransactionThroughput } from './get_transaction_throughput';
+
+export enum ApmTimeseriesType {
+  transactionThroughput = 'transaction_throughput',
+  transactionLatency = 'transaction_latency',
+  transactionFailureRate = 'transaction_failure_rate',
+  exitSpanThroughput = 'exit_span_throughput',
+  exitSpanLatency = 'exit_span_latency',
+  exitSpanFailureRate = 'exit_span_failure_rate',
+  errorEventRate = 'error_event_rate',
+}
+
+export const getApmTimeseriesRt = t.type({
+  stats: t.array(
+    t.intersection([
+      t.type({
+        'service.environment': environmentRt.props.environment,
+        'service.name': t.string,
+        title: t.string,
+        timeseries: t.union([
+          t.intersection([
+            t.type({
+              name: t.union([
+                t.literal(ApmTimeseriesType.transactionThroughput),
+                t.literal(ApmTimeseriesType.transactionFailureRate),
+              ]),
+            }),
+            t.partial({
+              'transaction.type': t.string,
+            }),
+          ]),
+          t.intersection([
+            t.type({
+              name: t.union([
+                t.literal(ApmTimeseriesType.exitSpanThroughput),
+                t.literal(ApmTimeseriesType.exitSpanFailureRate),
+                t.literal(ApmTimeseriesType.exitSpanLatency),
+              ]),
+            }),
+            t.partial({
+              'span.destination.service.resource': t.string,
+            }),
+          ]),
+          t.intersection([
+            t.type({
+              name: t.literal(ApmTimeseriesType.transactionLatency),
+              function: t.union([
+                t.literal(LatencyAggregationType.avg),
+                t.literal(LatencyAggregationType.p95),
+                t.literal(LatencyAggregationType.p99),
+              ]),
+            }),
+            t.partial({
+              'transaction.type': t.string,
+            }),
+          ]),
+          t.type({
+            name: t.literal(ApmTimeseriesType.errorEventRate),
+          }),
+        ]),
+      }),
+      t.partial({
+        filter: t.string,
+        offset: t.string,
+      }),
+    ])
+  ),
+  start: t.string,
+  end: t.string,
+});
+
+type ApmTimeseriesArgs = t.TypeOf<typeof getApmTimeseriesRt>;
+
+export interface ApmTimeseries {
+  stat: ApmTimeseriesArgs['stats'][number];
+  group: string;
+  id: string;
+  data: Array<{ x: number; y: number | null }>;
+  value: number | null;
+  start: number;
+  end: number;
+  unit: string;
+  changes: Array<{
+    change_point?: number | undefined;
+    r_value?: number | undefined;
+    trend?: string | undefined;
+    p_value: number;
+    date: string | undefined;
+    type: string;
+  }>;
+}
+
+export async function getApmTimeseries({
+  arguments: args,
+  apmEventClient,
+}: {
+  arguments: t.TypeOf<typeof getApmTimeseriesRt>;
+  apmEventClient: APMEventClient;
+}): Promise<ApmTimeseries[]> {
+  const start = datemath.parse(args.start)!.valueOf();
+  const end = datemath.parse(args.end)!.valueOf();
+
+  const { bucketSize, intervalString } = getBucketSize({
+    start,
+    end,
+    numBuckets: 100,
+  });
+
+  const sharedParameters = {
+    apmEventClient,
+    start,
+    end,
+    bucketSize,
+    intervalString,
+  };
+
+  return (
+    await Promise.all(
+      args.stats.map(async (stat) => {
+        const parameters = {
+          ...sharedParameters,
+          filter: [
+            ...rangeQuery(start, end),
+            ...termQuery(SERVICE_NAME, stat['service.name']),
+            ...kqlQuery(stat.filter),
+            ...environmentQuery(stat['service.environment']),
+          ],
+        };
+        const name = stat.timeseries.name;
+
+        async function fetchSeriesForStat() {
+          switch (name) {
+            case ApmTimeseriesType.transactionThroughput:
+              return await getTransactionThroughput({
+                ...parameters,
+                transactionType: stat.timeseries['transaction.type'],
+              });
+
+            case ApmTimeseriesType.transactionFailureRate:
+              return await getTransactionFailureRate({
+                ...parameters,
+                transactionType: stat.timeseries['transaction.type'],
+              });
+
+            case ApmTimeseriesType.transactionLatency:
+              return await getTransactionLatency({
+                ...parameters,
+                transactionType: stat.timeseries['transaction.type'],
+                latencyAggregationType: stat.timeseries.function,
+              });
+
+            case ApmTimeseriesType.exitSpanThroughput:
+              return await getExitSpanThroughput({
+                ...parameters,
+                spanDestinationServiceResource:
+                  stat.timeseries['span.destination.service.resource'],
+              });
+
+            case ApmTimeseriesType.exitSpanFailureRate:
+              return await getExitSpanFailureRate({
+                ...parameters,
+                spanDestinationServiceResource:
+                  stat.timeseries['span.destination.service.resource'],
+              });
+
+            case ApmTimeseriesType.exitSpanLatency:
+              return await getExitSpanLatency({
+                ...parameters,
+                spanDestinationServiceResource:
+                  stat.timeseries['span.destination.service.resource'],
+              });
+
+            case ApmTimeseriesType.errorEventRate:
+              return await getErrorEventRate(parameters);
+          }
+        }
+
+        const allFetchedSeries = await fetchSeriesForStat();
+        return allFetchedSeries.map((series) => ({ ...series, stat }));
+      })
+    )
+  ).flatMap((statResults) =>
+    statResults.flatMap((statResult) => {
+      const changePointType = Object.keys(
+        statResult.change_point?.type ?? {}
+      )?.[0];
+
+      return {
+        stat: statResult.stat,
+        group: statResult.stat.title,
+        id: statResult.groupBy,
+        data: statResult.data,
+        value: statResult.value,
+        start,
+        end,
+        unit: statResult.unit,
+        changes: [
+          ...(changePointType && changePointType !== 'indeterminable'
+            ? [
+                {
+                  date: statResult.change_point.bucket?.key,
+                  type: changePointType,
+                  ...statResult.change_point.type[changePointType],
+                },
+              ]
+            : []),
+        ],
+      };
+    })
+  );
+}

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

@@ -0,0 +1,193 @@
+/*
+ * 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 { ElasticsearchClient } from '@kbn/core/server';
+import * as t from 'io-ts';
+import { omit } from 'lodash';
+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 { createApmServerRoute } from '../apm_routes/create_apm_server_route';
+import {
+  CorrelationValue,
+  correlationValuesRouteRt,
+  getApmCorrelationValues,
+} from './get_apm_correlation_values';
+import {
+  type APMDownstreamDependency,
+  downstreamDependenciesRouteRt,
+  getAssistantDownstreamDependencies,
+} from './get_apm_downstream_dependencies';
+import { errorRouteRt, getApmErrorDocument } from './get_apm_error_document';
+import {
+  getApmServiceSummary,
+  type ServiceSummary,
+  serviceSummaryRouteRt,
+} from './get_apm_service_summary';
+import {
+  type ApmTimeseries,
+  getApmTimeseries,
+  getApmTimeseriesRt,
+} from './get_apm_timeseries';
+
+const getApmTimeSeriesRoute = createApmServerRoute({
+  endpoint: 'POST /internal/apm/assistant/get_apm_timeseries',
+  options: {
+    tags: ['access:apm', 'access:ai_assistant'],
+  },
+  params: t.type({
+    body: getApmTimeseriesRt,
+  }),
+  handler: async (
+    resources
+  ): Promise<{
+    content: Array<Omit<ApmTimeseries, 'data'>>;
+    data: ApmTimeseries[];
+  }> => {
+    const body = resources.params.body;
+
+    const apmEventClient = await getApmEventClient(resources);
+
+    const timeseries = await getApmTimeseries({
+      apmEventClient,
+      arguments: body,
+    });
+
+    return {
+      content: timeseries.map(
+        (series): Omit<ApmTimeseries, 'data'> => omit(series, 'data')
+      ),
+      data: timeseries,
+    };
+  },
+});
+
+const getApmServiceSummaryRoute = createApmServerRoute({
+  endpoint: 'GET /internal/apm/assistant/get_service_summary',
+  options: {
+    tags: ['access:apm', 'access:ai_assistant'],
+  },
+  params: t.type({
+    query: serviceSummaryRouteRt,
+  }),
+  handler: async (
+    resources
+  ): Promise<{
+    content: ServiceSummary;
+  }> => {
+    const args = resources.params.query;
+
+    const { context, request, plugins, logger } = resources;
+
+    const [
+      apmEventClient,
+      annotationsClient,
+      esClient,
+      apmAlertsClient,
+      mlClient,
+    ] = await Promise.all([
+      getApmEventClient(resources),
+      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,
+      }),
+    };
+  },
+});
+
+const getDownstreamDependenciesRoute = createApmServerRoute({
+  endpoint: 'GET /internal/apm/assistant/get_downstream_dependencies',
+  params: t.type({
+    query: downstreamDependenciesRouteRt,
+  }),
+  options: {
+    tags: ['access:apm'],
+  },
+  handler: async (
+    resources
+  ): Promise<{ content: APMDownstreamDependency[] }> => {
+    const { params } = resources;
+    const apmEventClient = await getApmEventClient(resources);
+    const { query } = params;
+
+    return {
+      content: await getAssistantDownstreamDependencies({
+        arguments: query,
+        apmEventClient,
+      }),
+    };
+  },
+});
+
+const getApmCorrelationValuesRoute = createApmServerRoute({
+  endpoint: 'POST /internal/apm/assistant/get_correlation_values',
+  params: t.type({
+    body: correlationValuesRouteRt,
+  }),
+  options: {
+    tags: ['access:apm'],
+  },
+  handler: async (resources): Promise<{ content: CorrelationValue[] }> => {
+    const { params } = resources;
+    const apmEventClient = await getApmEventClient(resources);
+    const { body } = params;
+
+    return {
+      content: await getApmCorrelationValues({
+        arguments: body,
+        apmEventClient,
+      }),
+    };
+  },
+});
+
+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: Partial<APMError> | undefined }> => {
+    const { params } = resources;
+    const apmEventClient = await getApmEventClient(resources);
+    const { query } = params;
+
+    return {
+      content: await getApmErrorDocument({
+        apmEventClient,
+        arguments: query,
+      }),
+    };
+  },
+});
+
+export const assistantRouteRepository = {
+  ...getApmTimeSeriesRoute,
+  ...getApmServiceSummaryRoute,
+  ...getApmErrorDocRoute,
+  ...getApmCorrelationValuesRoute,
+  ...getDownstreamDependenciesRoute,
+};

x-pack/plugins/apm/server/routes/typings.ts

@@ -39,6 +39,7 @@ export interface APMRouteCreateOptions {
       | 'access:ml:canGetJobs'
       | 'access:ml:canCreateJob'
       | 'access:ml:canCloseJob'
+      | 'access:ai_assistant'
     >;
     body?: { accepts: Array<'application/json' | 'multipart/form-data'> };
     disableTelemetry?: boolean;

x-pack/plugins/observability_ai_assistant/common/types.ts

@@ -5,7 +5,6 @@
  * 2.0.
  */
 
-import type { Serializable } from '@kbn/utility-types';
 import type { FromSchema } from 'json-schema-to-ts';
 import type { JSONSchema } from 'json-schema-to-ts';
 import React from 'react';
@@ -23,7 +22,6 @@ export interface Message {
   message: {
     content?: string;
     name?: string;
-    event?: string;
     role: MessageRole;
     function_call?: {
       name: string;
@@ -76,41 +74,44 @@ export interface ContextDefinition {
 }
 
 interface FunctionResponse {
-  content?: Serializable;
-  data?: Serializable;
+  content?: any;
+  data?: any;
 }
 
 interface FunctionOptions<TParameters extends CompatibleJSONSchema = CompatibleJSONSchema> {
   name: string;
   description: string;
+  descriptionForUser: string;
   parameters: TParameters;
   contexts: string[];
 }
 
-type RespondFunction<
-  TParameters extends CompatibleJSONSchema,
-  TResponse extends FunctionResponse
-> = (options: { arguments: FromSchema<TParameters> }, signal: AbortSignal) => Promise<TResponse>;
+type RespondFunction<TArguments, TResponse extends FunctionResponse> = (
+  options: { arguments: TArguments },
+  signal: AbortSignal
+) => Promise<TResponse>;
 
-type RenderFunction<TResponse extends FunctionResponse> = (options: {
+type RenderFunction<TArguments, TResponse extends FunctionResponse> = (options: {
+  arguments: TArguments;
   response: TResponse;
 }) => React.ReactNode;
 
 export interface FunctionDefinition {
   options: FunctionOptions;
   respond: (options: { arguments: any }, signal: AbortSignal) => Promise<FunctionResponse>;
-  render?: RenderFunction<any>;
+  render?: RenderFunction<any, any>;
 }
 
 export type RegisterContextDefinition = (options: ContextDefinition) => void;
 
 export type RegisterFunctionDefinition = <
   TParameters extends CompatibleJSONSchema,
-  TResponse extends FunctionResponse
+  TResponse extends FunctionResponse,
+  TArguments = FromSchema<TParameters>
 >(
   options: FunctionOptions<TParameters>,
-  respond: RespondFunction<TParameters, TResponse>,
-  render?: RenderFunction<TResponse>
+  respond: RespondFunction<TArguments, TResponse>,
+  render?: RenderFunction<TArguments, TResponse>
 ) => void;
 
 export type ContextRegistry = Map<string, ContextDefinition>;

x-pack/plugins/observability_ai_assistant/public/application.tsx

@@ -6,12 +6,13 @@
  */
 import { EuiErrorBoundary } from '@elastic/eui';
 import type { CoreStart, CoreTheme } from '@kbn/core/public';
-import { KibanaContextProvider, KibanaThemeProvider } from '@kbn/kibana-react-plugin/public';
+import { KibanaContextProvider } from '@kbn/kibana-react-plugin/public';
 import { RedirectAppLinks } from '@kbn/shared-ux-link-redirect-app';
 import { RouteRenderer, RouterProvider } from '@kbn/typed-react-router-config';
 import type { History } from 'history';
-import React from 'react';
+import React, { useMemo } from 'react';
 import type { Observable } from 'rxjs';
+import { KibanaThemeProvider } from '@kbn/react-kibana-context-theme';
 import { ObservabilityAIAssistantProvider } from './context/observability_ai_assistant_provider';
 import { observabilityAIAssistantRouter } from './routes/config';
 import type {
@@ -32,9 +33,12 @@ export function Application({
   pluginsStart: ObservabilityAIAssistantPluginStartDependencies;
   service: ObservabilityAIAssistantService;
 }) {
+  const theme = useMemo(() => {
+    return { theme$ };
+  }, [theme$]);
   return (
     <EuiErrorBoundary>
-      <KibanaThemeProvider theme$={theme$}>
+      <KibanaThemeProvider theme={theme}>
         <KibanaContextProvider
           services={{
             ...coreStart,

x-pack/plugins/observability_ai_assistant/public/components/action_menu_item/action_menu_item.tsx

@@ -4,9 +4,11 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import { EuiFlexGroup, EuiFlexItem, EuiHeaderLink } from '@elastic/eui';
+import { EuiFlexGroup, EuiFlexItem, EuiHeaderLink, EuiLoadingSpinner } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 import React, { useState } from 'react';
+import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
+import { useAbortableAsync } from '../../hooks/use_abortable_async';
 import { useConversation } from '../../hooks/use_conversation';
 import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import { EMPTY_CONVERSATION_TITLE } from '../../i18n';
@@ -16,12 +18,23 @@ import { ChatFlyout } from '../chat/chat_flyout';
 export function ObservabilityAIAssistantActionMenuItem() {
   const service = useObservabilityAIAssistant();
 
+  const [isOpen, setIsOpen] = useState(false);
+
+  const chatService = useAbortableAsync(
+    ({ signal }) => {
+      if (!isOpen) {
+        return Promise.resolve(undefined);
+      }
+      return service.start({ signal });
+    },
+    [service, isOpen]
+  );
+
   const [conversationId, setConversationId] = useState<string>();
 
-  const { conversation, displayedMessages, setDisplayedMessages, save } =
-    useConversation(conversationId);
-
-  const [isOpen, setIsOpen] = useState(false);
+  const { conversation, displayedMessages, setDisplayedMessages, save } = useConversation({
+    conversationId,
+  });
 
   if (!service.isEnabled()) {
     return null;
@@ -37,7 +50,11 @@ export function ObservabilityAIAssistantActionMenuItem() {
       >
         <EuiFlexGroup gutterSize="s" alignItems="center">
           <EuiFlexItem grow={false}>
-            <AssistantAvatar size="xs" />
+            {!isOpen || chatService.value ? (
+              <AssistantAvatar size="xs" />
+            ) : (
+              <EuiLoadingSpinner size="s" />
+            )}
           </EuiFlexItem>
           <EuiFlexItem grow={false}>
             {i18n.translate('xpack.observabilityAiAssistant.actionMenuItemLabel', {
@@ -46,25 +63,29 @@ export function ObservabilityAIAssistantActionMenuItem() {
           </EuiFlexItem>
         </EuiFlexGroup>
       </EuiHeaderLink>
-      <ChatFlyout
-        isOpen={isOpen}
-        title={conversation.value?.conversation.title ?? EMPTY_CONVERSATION_TITLE}
-        messages={displayedMessages}
-        conversationId={conversationId}
-        onClose={() => {
-          setIsOpen(() => false);
-        }}
-        onChatComplete={(messages) => {
-          save(messages)
-            .then((nextConversation) => {
-              setConversationId(nextConversation.conversation.id);
-            })
-            .catch(() => {});
-        }}
-        onChatUpdate={(nextMessages) => {
-          setDisplayedMessages(nextMessages);
-        }}
-      />
+      {chatService.value ? (
+        <ObservabilityAIAssistantChatServiceProvider value={chatService.value}>
+          <ChatFlyout
+            isOpen={isOpen}
+            title={conversation.value?.conversation.title ?? EMPTY_CONVERSATION_TITLE}
+            messages={displayedMessages}
+            conversationId={conversationId}
+            onClose={() => {
+              setIsOpen(() => false);
+            }}
+            onChatComplete={(messages) => {
+              save(messages)
+                .then((nextConversation) => {
+                  setConversationId(nextConversation.conversation.id);
+                })
+                .catch(() => {});
+            }}
+            onChatUpdate={(nextMessages) => {
+              setDisplayedMessages(nextMessages);
+            }}
+          />
+        </ObservabilityAIAssistantChatServiceProvider>
+      ) : null}
     </>
   );
 }

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

@@ -18,12 +18,14 @@ import type { AuthenticatedUser } from '@kbn/security-plugin/common';
 import React from 'react';
 import type { Message } from '../../../common/types';
 import type { UseGenAIConnectorsResult } from '../../hooks/use_genai_connectors';
+import type { UseKnowledgeBaseResult } from '../../hooks/use_knowledge_base';
 import { useTimeline } from '../../hooks/use_timeline';
-import { ObservabilityAIAssistantService } from '../../types';
+import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
 import { MissingCredentialsCallout } from '../missing_credentials_callout';
 import { ChatHeader } from './chat_header';
 import { ChatPromptEditor } from './chat_prompt_editor';
 import { ChatTimeline } from './chat_timeline';
+import { KnowledgeBaseCallout } from './knowledge_base_callout';
 
 const containerClassName = css`
   max-height: 100%;
@@ -42,8 +44,8 @@ export function ChatBody({
   title,
   messages,
   connectors,
+  knowledgeBase,
   currentUser,
-  service,
   connectorsManagementHref,
   onChatUpdate,
   onChatComplete,
@@ -51,34 +53,36 @@ export function ChatBody({
   title: string;
   messages: Message[];
   connectors: UseGenAIConnectorsResult;
+  knowledgeBase: UseKnowledgeBaseResult;
   currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
-  service: ObservabilityAIAssistantService;
   connectorsManagementHref: string;
   onChatUpdate: (messages: Message[]) => void;
   onChatComplete: (messages: Message[]) => void;
 }) {
+  const chatService = useObservabilityAIAssistantChatService();
+
   const timeline = useTimeline({
     messages,
     connectors,
     currentUser,
-    service,
+    chatService,
     onChatUpdate,
     onChatComplete,
   });
 
   let footer: React.ReactNode;
 
-  if (connectors.loading || connectors.connectors?.length === 0) {
+  if (connectors.loading || knowledgeBase.status.loading) {
+    footer = (
+      <EuiFlexItem className={loadingSpinnerContainerClassName}>
+        <EuiLoadingSpinner />
+      </EuiFlexItem>
+    );
+  } else if (connectors.connectors?.length === 0) {
     footer = (
       <>
         <EuiSpacer size="l" />
-        {connectors.connectors?.length === 0 ? (
-          <MissingCredentialsCallout connectorsManagementHref={connectorsManagementHref} />
-        ) : (
-          <EuiFlexItem className={loadingSpinnerContainerClassName}>
-            <EuiLoadingSpinner />
-          </EuiFlexItem>
-        )}
+        <MissingCredentialsCallout connectorsManagementHref={connectorsManagementHref} />
       </>
     );
   } else {
@@ -118,6 +122,9 @@ export function ChatBody({
           <ChatHeader title={title} connectors={connectors} />
         </EuiPanel>
       </EuiFlexItem>
+      <EuiFlexItem grow={false}>
+        <KnowledgeBaseCallout knowledgeBase={knowledgeBase} />
+      </EuiFlexItem>
       <EuiFlexItem grow={false}>
         <EuiHorizontalRule margin="none" />
       </EuiFlexItem>

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx

@@ -7,7 +7,7 @@
 
 import { ComponentStory } from '@storybook/react';
 import React from 'react';
-import { getSystemMessage } from '../../service/get_system_message';
+import { getAssistantSetupMessage } from '../../service/get_assistant_setup_message';
 import { KibanaReactStorybookDecorator } from '../../utils/storybook_decorator';
 import { ChatFlyout as Component } from './chat_flyout';
 
@@ -30,7 +30,7 @@ const Template: ComponentStory<typeof Component> = (props: ChatFlyoutProps) => {
 const defaultProps: ChatFlyoutProps = {
   isOpen: true,
   title: 'How is this working',
-  messages: [getSystemMessage()],
+  messages: [getAssistantSetupMessage({ contexts: [] })],
   onClose: () => {},
 };
 

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.tsx

@@ -6,13 +6,13 @@
  */
 import { EuiFlexGroup, EuiFlexItem, EuiFlyout, EuiLink, EuiPanel, useEuiTheme } from '@elastic/eui';
 import { css } from '@emotion/css';
-import React from 'react';
 import { i18n } from '@kbn/i18n';
+import React from 'react';
 import type { Message } from '../../../common/types';
 import { useCurrentUser } from '../../hooks/use_current_user';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
 import { useKibana } from '../../hooks/use_kibana';
-import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
+import { useKnowledgeBase } from '../../hooks/use_knowledge_base';
 import { useObservabilityAIAssistantRouter } from '../../hooks/use_observability_ai_assistant_router';
 import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
 import { ChatBody } from './chat_body';
@@ -21,6 +21,10 @@ const containerClassName = css`
   max-height: 100%;
 `;
 
+const bodyClassName = css`
+  overflow-y: auto;
+`;
+
 export function ChatFlyout({
   title,
   messages,
@@ -46,12 +50,12 @@ export function ChatFlyout({
     services: { http },
   } = useKibana();
 
-  const service = useObservabilityAIAssistant();
-
   const { euiTheme } = useEuiTheme();
 
   const router = useObservabilityAIAssistantRouter();
 
+  const knowledgeBase = useKnowledgeBase();
+
   return isOpen ? (
     <EuiFlyout onClose={onClose}>
       <EuiFlexGroup
@@ -86,14 +90,14 @@ export function ChatFlyout({
             )}
           </EuiPanel>
         </EuiFlexItem>
-        <EuiFlexItem>
+        <EuiFlexItem grow className={bodyClassName}>
           <ChatBody
-            service={service}
             connectors={connectors}
             title={title}
             messages={messages}
             currentUser={currentUser}
             connectorsManagementHref={getConnectorsManagementHref(http)}
+            knowledgeBase={knowledgeBase}
             onChatUpdate={(nextMessages) => {
               if (onChatUpdate) {
                 onChatUpdate(nextMessages);

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item.tsx

@@ -5,191 +5,193 @@
  * 2.0.
  */
 
-import {
-  EuiButtonIcon,
-  EuiComment,
-  EuiContextMenuItem,
-  EuiContextMenuPanel,
-  EuiFlexGroup,
-  EuiFlexItem,
-  EuiPopover,
-} from '@elastic/eui';
-import { css } from '@emotion/css';
-import { i18n } from '@kbn/i18n';
 import React, { useState } from 'react';
-import { MessageRole } from '../../../common/types';
-import { Feedback, FeedbackButtons } from '../feedback_buttons';
-import { MessagePanel } from '../message_panel/message_panel';
-import { MessageText } from '../message_panel/message_text';
-import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
-import { StopGeneratingButton } from '../buttons/stop_generating_button';
+import { css } from '@emotion/css';
+import {
+  EuiAccordion,
+  EuiComment,
+  EuiErrorBoundary,
+  EuiPanel,
+  EuiSpacer,
+  useGeneratedHtmlId,
+} from '@elastic/eui';
+import { ChatItemActions } from './chat_item_actions';
 import { ChatItemAvatar } from './chat_item_avatar';
-import { ChatItemTitle } from './chat_item_title';
+import { ChatItemContentInlinePromptEditor } from './chat_item_content_inline_prompt_editor';
+import { ChatItemControls } from './chat_item_controls';
 import { ChatTimelineItem } from './chat_timeline';
-
-export interface ChatItemAction {
-  id: string;
-  label: string;
-  icon?: string;
-  handler: () => void;
-}
+import { getRoleTranslation } from '../../utils/get_role_translation';
+import type { Feedback } from '../feedback_buttons';
+import { Message } from '../../../common';
 
 export interface ChatItemProps extends ChatTimelineItem {
-  onEditSubmit: (content: string) => void;
+  onEditSubmit: (message: Message) => Promise<void>;
   onFeedbackClick: (feedback: Feedback) => void;
   onRegenerateClick: () => void;
   onStopGeneratingClick: () => void;
 }
 
-const euiCommentClassName = css`
-  .euiCommentEvent__headerEvent {
-    flex-grow: 1;
+const normalMessageClassName = css`
+  .euiCommentEvent__header {
+    padding: 4px 8px;
   }
 
-  > div:last-child {
-    overflow: hidden;
+  .euiCommentEvent__body {
+    padding: 0;
+  }
+  /* targets .*euiTimelineItemEvent-top, makes sure text properly wraps and doesn't overflow */
+  > :last-child {
+    overflow-x: hidden;
+  }
+`;
+
+const noPanelMessageClassName = css`
+  .euiCommentEvent {
+    border: none;
+  }
+
+  .euiCommentEvent__header {
+    background: transparent;
+    border-block-end: none;
+  }
+
+  .euiCommentEvent__body {
+    display: none;
+  }
+`;
+
+const accordionButtonClassName = css`
+  .euiAccordion__iconButton {
+    display: none;
   }
 `;
 
 export function ChatItem({
-  title,
+  actions: { canCopy, canEdit, canGiveFeedback, canRegenerate },
+  display: { collapsed },
   content,
-  canEdit,
-  canGiveFeedback,
-  canRegenerate,
-  role,
-  loading,
-  error,
   currentUser,
+  element,
+  error,
+  function_call: functionCall,
+  loading,
+  role,
+  title,
   onEditSubmit,
+  onFeedbackClick,
   onRegenerateClick,
   onStopGeneratingClick,
-  onFeedbackClick,
 }: ChatItemProps) {
-  const [isActionsPopoverOpen, setIsActionsPopover] = useState(false);
+  const accordionId = useGeneratedHtmlId({ prefix: 'chat' });
 
-  const handleClickActions = () => {
-    setIsActionsPopover(!isActionsPopoverOpen);
+  const [editing, setEditing] = useState<boolean>(false);
+  const [expanded, setExpanded] = useState<boolean>(Boolean(element));
+
+  const actions = [canCopy, collapsed, canCopy].filter(Boolean);
+
+  const noBodyMessageClassName = css`
+    .euiCommentEvent__header {
+      padding: 4px 8px;
+    }
+
+    .euiCommentEvent__body {
+      padding: 0;
+      height: ${expanded ? 'fit-content' : '0px'};
+      overflow: hidden;
+    }
+  `;
+
+  const handleToggleExpand = () => {
+    setExpanded(!expanded);
+
+    if (editing) {
+      setEditing(false);
+    }
   };
 
-  const [_, setEditing] = useState(false);
+  const handleToggleEdit = () => {
+    if (collapsed && !expanded) {
+      setExpanded(true);
+    }
+    setEditing(!editing);
+  };
 
-  const actions: ChatItemAction[] = canEdit
-    ? [
-        {
-          id: 'edit',
-          label: i18n.translate('xpack.observabilityAiAssistant.chatTimeline.actions.editMessage', {
-            defaultMessage: 'Edit message',
-          }),
-          handler: () => {
-            setEditing(false);
-            setIsActionsPopover(false);
-          },
-        },
-      ]
-    : [];
+  const handleInlineEditSubmit = (message: Message) => {
+    handleToggleEdit();
+    return onEditSubmit(message);
+  };
 
-  let controls: React.ReactNode;
+  const handleCopyToClipboard = () => {
+    navigator.clipboard.writeText(content || '');
+  };
 
-  const displayFeedback = !error && canGiveFeedback;
-  const displayRegenerate = !loading && canRegenerate;
+  let contentElement: React.ReactNode =
+    content || error ? (
+      <ChatItemContentInlinePromptEditor
+        content={content}
+        editing={editing}
+        functionCall={functionCall}
+        loading={loading}
+        onSubmit={handleInlineEditSubmit}
+      />
+    ) : null;
 
-  if (loading) {
-    controls = <StopGeneratingButton onClick={onStopGeneratingClick} />;
-  } else if (displayFeedback || displayRegenerate) {
-    controls = (
-      <EuiFlexGroup justifyContent="flexEnd">
-        {displayFeedback ? (
-          <EuiFlexItem grow={true}>
-            <FeedbackButtons onClickFeedback={onFeedbackClick} />
-          </EuiFlexItem>
-        ) : null}
-        {displayRegenerate ? (
-          <EuiFlexItem grow={false} style={{ alignSelf: 'flex-end' }}>
-            <RegenerateResponseButton onClick={onRegenerateClick} />
-          </EuiFlexItem>
-        ) : null}
-      </EuiFlexGroup>
+  if (collapsed) {
+    contentElement = (
+      <EuiAccordion
+        id={accordionId}
+        className={accordionButtonClassName}
+        forceState={expanded ? 'open' : 'closed'}
+        onToggle={handleToggleExpand}
+      >
+        <EuiSpacer size="s" />
+        {contentElement}
+      </EuiAccordion>
     );
   }
 
   return (
     <EuiComment
-      event={
-        <ChatItemTitle
-          actionsTrigger={
-            actions.length ? (
-              <EuiPopover
-                anchorPosition="downLeft"
-                button={
-                  <EuiButtonIcon
-                    aria-label={i18n.translate(
-                      'xpack.observabilityAiAssistant.chatTimeline.actions',
-                      {
-                        defaultMessage: 'Actions',
-                      }
-                    )}
-                    color="text"
-                    display="empty"
-                    iconType="boxesHorizontal"
-                    size="s"
-                    onClick={handleClickActions}
-                  />
-                }
-                panelPaddingSize="s"
-                closePopover={handleClickActions}
-                isOpen={isActionsPopoverOpen}
-              >
-                <EuiContextMenuPanel
-                  size="s"
-                  items={actions.map(({ id, icon, label, handler }) => (
-                    <EuiContextMenuItem key={id} icon={icon} onClick={handler}>
-                      {label}
-                    </EuiContextMenuItem>
-                  ))}
-                />
-              </EuiPopover>
-            ) : null
-          }
-          title={title}
-        />
-      }
-      className={euiCommentClassName}
       timelineAvatar={
         <ChatItemAvatar loading={loading && !content} currentUser={currentUser} role={role} />
       }
       username={getRoleTranslation(role)}
-    >
-      {content || error || controls ? (
-        <MessagePanel
-          body={
-            content || loading ? <MessageText content={content || ''} loading={loading} /> : null
-          }
-          error={error}
-          controls={controls}
+      event={title}
+      actions={
+        <ChatItemActions
+          canCopy={canCopy}
+          canEdit={canEdit}
+          collapsed={collapsed}
+          editing={editing}
+          expanded={expanded}
+          onCopyToClipboard={handleCopyToClipboard}
+          onToggleEdit={handleToggleEdit}
+          onToggleExpand={handleToggleExpand}
         />
-      ) : null}
+      }
+      className={
+        actions.length === 0 && !content
+          ? noPanelMessageClassName
+          : collapsed
+          ? noBodyMessageClassName
+          : normalMessageClassName
+      }
+    >
+      <EuiPanel hasShadow={false} paddingSize="s">
+        {element ? <EuiErrorBoundary>{element}</EuiErrorBoundary> : null}
+
+        {contentElement}
+      </EuiPanel>
+
+      <ChatItemControls
+        canGiveFeedback={canGiveFeedback}
+        canRegenerate={canRegenerate}
+        error={error}
+        loading={loading}
+        onFeedbackClick={onFeedbackClick}
+        onRegenerateClick={onRegenerateClick}
+        onStopGeneratingClick={onStopGeneratingClick}
+      />
     </EuiComment>
   );
 }
-
-const getRoleTranslation = (role: MessageRole) => {
-  if (role === MessageRole.User) {
-    return i18n.translate('xpack.observabilityAiAssistant.chatTimeline.messages.user.label', {
-      defaultMessage: 'You',
-    });
-  }
-
-  if (role === MessageRole.System) {
-    return i18n.translate('xpack.observabilityAiAssistant.chatTimeline.messages.system.label', {
-      defaultMessage: 'System',
-    });
-  }
-
-  return i18n.translate(
-    'xpack.observabilityAiAssistant.chatTimeline.messages.elasticAssistant.label',
-    {
-      defaultMessage: 'Elastic Assistant',
-    }
-  );
-};

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_actions.tsx

@@ -0,0 +1,114 @@
+/*
+ * 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 React, { useEffect, useState } from 'react';
+import { i18n } from '@kbn/i18n';
+import { EuiButtonIcon, EuiPopover, EuiText } from '@elastic/eui';
+
+export function ChatItemActions({
+  canCopy,
+  canEdit,
+  collapsed,
+  editing,
+  expanded,
+  onToggleEdit,
+  onToggleExpand,
+  onCopyToClipboard,
+}: {
+  canCopy: boolean;
+  canEdit: boolean;
+  collapsed: boolean;
+  editing: boolean;
+  expanded: boolean;
+  onToggleEdit: () => void;
+  onToggleExpand: () => void;
+  onCopyToClipboard: () => void;
+}) {
+  const [isPopoverOpen, setIsPopoverOpen] = useState<string | undefined>();
+
+  useEffect(() => {
+    const timeout = setTimeout(() => {
+      if (isPopoverOpen) {
+        setIsPopoverOpen(undefined);
+      }
+    }, 800);
+
+    return () => {
+      clearTimeout(timeout);
+    };
+  }, [isPopoverOpen]);
+
+  return (
+    <>
+      {canEdit ? (
+        <EuiButtonIcon
+          aria-label={i18n.translate(
+            'xpack.observabilityAiAssistant.chatTimeline.actions.editPrompt',
+            {
+              defaultMessage: 'Edit prompt',
+            }
+          )}
+          color="text"
+          display={editing ? 'fill' : 'empty'}
+          iconType="documentEdit"
+          onClick={onToggleEdit}
+        />
+      ) : null}
+
+      {collapsed ? (
+        <EuiButtonIcon
+          aria-label={i18n.translate(
+            'xpack.observabilityAiAssistant.chatTimeline.actions.inspectPrompt',
+            {
+              defaultMessage: 'Inspect prompt',
+            }
+          )}
+          color="text"
+          display={expanded ? 'fill' : 'empty'}
+          iconType={expanded ? 'eyeClosed' : 'eye'}
+          onClick={onToggleExpand}
+        />
+      ) : null}
+
+      {canCopy ? (
+        <EuiPopover
+          button={
+            <EuiButtonIcon
+              aria-label={i18n.translate(
+                'xpack.observabilityAiAssistant.chatTimeline.actions.copyMessage',
+                {
+                  defaultMessage: 'Copy message',
+                }
+              )}
+              color="text"
+              iconType="copyClipboard"
+              display={isPopoverOpen === 'copy' ? 'fill' : 'empty'}
+              onClick={() => {
+                setIsPopoverOpen('copy');
+                onCopyToClipboard();
+              }}
+            />
+          }
+          isOpen={isPopoverOpen === 'copy'}
+          panelPaddingSize="s"
+          closePopover={() => setIsPopoverOpen(undefined)}
+        >
+          <EuiText size="s">
+            <p>
+              {i18n.translate(
+                'xpack.observabilityAiAssistant.chatTimeline.actions.copyMessageSuccessful',
+                {
+                  defaultMessage: 'Copied message',
+                }
+              )}
+            </p>
+          </EuiText>
+        </EuiPopover>
+      ) : null}
+    </>
+  );
+}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx

@@ -0,0 +1,46 @@
+/*
+ * 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 React from 'react';
+import { MessageText } from '../message_panel/message_text';
+import { ChatPromptEditor } from './chat_prompt_editor';
+import { MessageRole, type Message } from '../../../common';
+
+interface Props {
+  content: string | undefined;
+  functionCall:
+    | {
+        name: string;
+        arguments?: string | undefined;
+        trigger: MessageRole;
+      }
+    | undefined;
+  loading: boolean;
+  editing: boolean;
+  onSubmit: (message: Message) => Promise<void>;
+}
+export function ChatItemContentInlinePromptEditor({
+  content,
+  functionCall,
+  editing,
+  loading,
+  onSubmit,
+}: Props) {
+  return !editing ? (
+    <MessageText content={content || ''} loading={loading} />
+  ) : (
+    <ChatPromptEditor
+      disabled={false}
+      loading={false}
+      initialPrompt={content}
+      initialFunctionPayload={functionCall?.arguments}
+      initialSelectedFunctionName={functionCall?.name}
+      trigger={functionCall?.trigger}
+      onSubmit={onSubmit}
+    />
+  );
+}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_controls.tsx

@@ -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 React from 'react';
+import {
+  EuiFlexGroup,
+  EuiFlexItem,
+  EuiHorizontalRule,
+  EuiPanel,
+  EuiSpacer,
+  useEuiTheme,
+} from '@elastic/eui';
+import { Feedback, FeedbackButtons } from '../feedback_buttons';
+import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
+import { StopGeneratingButton } from '../buttons/stop_generating_button';
+
+export function ChatItemControls({
+  error,
+  loading,
+  canRegenerate,
+  canGiveFeedback,
+  onFeedbackClick,
+  onRegenerateClick,
+  onStopGeneratingClick,
+}: {
+  error: any;
+  loading: boolean;
+  canRegenerate: boolean;
+  canGiveFeedback: boolean;
+  onFeedbackClick: (feedback: Feedback) => void;
+  onRegenerateClick: () => void;
+  onStopGeneratingClick: () => void;
+}) {
+  const { euiTheme } = useEuiTheme();
+
+  const displayFeedback = !error && canGiveFeedback;
+  const displayRegenerate = !loading && canRegenerate;
+
+  let controls;
+
+  if (loading) {
+    controls = <StopGeneratingButton onClick={onStopGeneratingClick} />;
+  } else if (displayFeedback || displayRegenerate) {
+    controls = (
+      <EuiFlexGroup justifyContent="flexEnd">
+        {displayFeedback ? (
+          <EuiFlexItem grow={true}>
+            <FeedbackButtons onClickFeedback={onFeedbackClick} />
+          </EuiFlexItem>
+        ) : null}
+        {displayRegenerate ? (
+          <EuiFlexItem grow={false} style={{ alignSelf: 'flex-end' }}>
+            <RegenerateResponseButton onClick={onRegenerateClick} />
+          </EuiFlexItem>
+        ) : null}
+      </EuiFlexGroup>
+    );
+  } else {
+    controls = null;
+  }
+
+  return controls ? (
+    <>
+      <EuiSpacer size="s" />
+      <EuiHorizontalRule margin="none" color={euiTheme.colors.lightestShade} />
+      <EuiPanel hasShadow={false} paddingSize="s">
+        {controls}
+      </EuiPanel>
+    </>
+  ) : null;
+}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx

@@ -5,48 +5,65 @@
  * 2.0.
  */
 
-import React, { useCallback, useEffect, useRef, useState } from 'react';
 import {
-  EuiButtonIcon,
   EuiButtonEmpty,
-  EuiFieldText,
+  EuiButtonIcon,
   EuiFlexGroup,
   EuiFlexItem,
-  EuiSpacer,
   EuiPanel,
+  EuiSpacer,
+  EuiTextArea,
   keys,
+  EuiFocusTrap,
 } from '@elastic/eui';
-import { CodeEditor } from '@kbn/kibana-react-plugin/public';
 import { i18n } from '@kbn/i18n';
-import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
+import { CodeEditor } from '@kbn/kibana-react-plugin/public';
+import React, { useCallback, useEffect, useRef, useState } from 'react';
+import { MessageRole, type Message } from '../../../common';
 import { useJsonEditorModel } from '../../hooks/use_json_editor_model';
-import { type Message, MessageRole } from '../../../common';
-import type { FunctionDefinition } from '../../../common/types';
 import { FunctionListPopover } from './function_list_popover';
 
 export interface ChatPromptEditorProps {
   disabled: boolean;
   loading: boolean;
+  initialPrompt?: string;
+  initialSelectedFunctionName?: string;
+  initialFunctionPayload?: string;
+  trigger?: MessageRole;
   onSubmit: (message: Message) => Promise<void>;
 }
 
-export function ChatPromptEditor({ onSubmit, disabled, loading }: ChatPromptEditorProps) {
-  const { getFunctions } = useObservabilityAIAssistant();
-  const functions = getFunctions();
+export function ChatPromptEditor({
+  disabled,
+  loading,
+  initialPrompt,
+  initialSelectedFunctionName,
+  initialFunctionPayload,
+  onSubmit,
+}: ChatPromptEditorProps) {
+  const isFocusTrapEnabled = Boolean(initialPrompt);
 
-  const [prompt, setPrompt] = useState('');
-  const [functionPayload, setFunctionPayload] = useState<string | undefined>('');
-  const [selectedFunction, setSelectedFunction] = useState<FunctionDefinition | undefined>();
+  const [prompt, setPrompt] = useState(initialPrompt);
 
-  const { model, initialJsonString } = useJsonEditorModel(selectedFunction);
+  const [selectedFunctionName, setSelectedFunctionName] = useState<string | undefined>(
+    initialSelectedFunctionName
+  );
+  const [functionPayload, setFunctionPayload] = useState<string | undefined>(
+    initialFunctionPayload
+  );
 
-  const ref = useRef<HTMLInputElement>(null);
+  const { model, initialJsonString } = useJsonEditorModel({
+    functionName: selectedFunctionName,
+    initialJson: initialFunctionPayload,
+  });
+
+  const textAreaRef = useRef<HTMLTextAreaElement>(null);
 
   useEffect(() => {
     setFunctionPayload(initialJsonString);
-  }, [initialJsonString, selectedFunction]);
+  }, [initialJsonString, selectedFunctionName]);
 
-  const handleChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+  const handleChange = (event: React.ChangeEvent<HTMLTextAreaElement>) => {
     setPrompt(event.currentTarget.value);
   };
 
@@ -55,8 +72,22 @@ export function ChatPromptEditor({ onSubmit, disabled, loading }: ChatPromptEdit
   };
 
   const handleClearSelection = () => {
-    setSelectedFunction(undefined);
+    setSelectedFunctionName(undefined);
     setFunctionPayload('');
+    setPrompt('');
+  };
+
+  const handleSelectFunction = (functionName: string) => {
+    setPrompt('');
+    setFunctionPayload('');
+    setSelectedFunctionName(functionName);
+  };
+
+  const handleResizeTextArea = () => {
+    if (textAreaRef.current) {
+      textAreaRef.current.style.height = 'auto';
+      textAreaRef.current.style.height = textAreaRef.current?.scrollHeight + 'px';
+    }
   };
 
   const handleSubmit = useCallback(async () => {
@@ -65,20 +96,25 @@ export function ChatPromptEditor({ onSubmit, disabled, loading }: ChatPromptEdit
 
     setPrompt('');
     setFunctionPayload(undefined);
+    handleResizeTextArea();
 
     try {
-      if (selectedFunction) {
+      if (selectedFunctionName) {
         await onSubmit({
           '@timestamp': new Date().toISOString(),
           message: {
-            role: MessageRole.Function,
+            role: MessageRole.Assistant,
+            content: '',
             function_call: {
-              name: selectedFunction.options.name,
+              name: selectedFunctionName,
               trigger: MessageRole.User,
               arguments: currentPayload,
             },
           },
         });
+
+        setFunctionPayload(undefined);
+        setSelectedFunctionName(undefined);
       } else {
         await onSubmit({
           '@timestamp': new Date().toISOString(),
@@ -89,131 +125,144 @@ export function ChatPromptEditor({ onSubmit, disabled, loading }: ChatPromptEdit
     } catch (_) {
       setPrompt(currentPrompt);
     }
-  }, [functionPayload, onSubmit, prompt, selectedFunction]);
+  }, [functionPayload, onSubmit, prompt, selectedFunctionName]);
 
   useEffect(() => {
     const keyboardListener = (event: KeyboardEvent) => {
-      if (event.key === keys.ENTER) {
+      if (!event.shiftKey && event.key === keys.ENTER && (prompt || selectedFunctionName)) {
+        event.preventDefault();
         handleSubmit();
       }
     };
 
-    window.addEventListener('keyup', keyboardListener);
+    window.addEventListener('keypress', keyboardListener);
 
     return () => {
-      window.removeEventListener('keyup', keyboardListener);
+      window.removeEventListener('keypress', keyboardListener);
     };
-  }, [handleSubmit]);
+  }, [handleSubmit, prompt, selectedFunctionName]);
 
   useEffect(() => {
-    if (ref.current) {
-      ref.current.focus();
+    const textarea = textAreaRef.current;
+
+    if (textarea) {
+      textarea.focus();
+      textarea.addEventListener('input', handleResizeTextArea, false);
     }
+
+    return () => {
+      textarea?.removeEventListener('input', handleResizeTextArea, false);
+    };
   });
 
   return (
-    <EuiFlexGroup gutterSize="s" responsive={false}>
-      <EuiFlexItem grow>
-        <EuiFlexGroup direction="column" gutterSize="s">
-          <EuiFlexItem>
-            <EuiFlexGroup responsive={false}>
-              <EuiFlexItem grow>
-                <FunctionListPopover
-                  functions={functions}
-                  selectedFunction={selectedFunction}
-                  onSelectFunction={setSelectedFunction}
-                />
-              </EuiFlexItem>
-              <EuiFlexItem grow={false}>
-                {selectedFunction ? (
-                  <EuiButtonEmpty
-                    iconType="cross"
-                    iconSide="right"
-                    size="xs"
-                    onClick={handleClearSelection}
-                  >
-                    {i18n.translate('xpack.observabilityAiAssistant.prompt.emptySelection', {
-                      defaultMessage: 'Empty selection',
-                    })}
-                  </EuiButtonEmpty>
-                ) : null}
-              </EuiFlexItem>
-            </EuiFlexGroup>
-          </EuiFlexItem>
-          <EuiFlexItem>
-            {selectedFunction ? (
-              <EuiPanel borderRadius="none" color="subdued" hasShadow={false} paddingSize="xs">
-                <CodeEditor
-                  aria-label="payloadEditor"
+    <EuiFocusTrap disabled={!isFocusTrapEnabled}>
+      <EuiFlexGroup gutterSize="s" responsive={false}>
+        <EuiFlexItem grow>
+          <EuiFlexGroup direction="column" gutterSize="s">
+            <EuiFlexItem>
+              <EuiFlexGroup responsive={false}>
+                <EuiFlexItem grow>
+                  <FunctionListPopover
+                    selectedFunctionName={selectedFunctionName}
+                    onSelectFunction={handleSelectFunction}
+                  />
+                </EuiFlexItem>
+                <EuiFlexItem grow={false}>
+                  {selectedFunctionName ? (
+                    <EuiButtonEmpty
+                      iconType="cross"
+                      iconSide="right"
+                      size="xs"
+                      onClick={handleClearSelection}
+                    >
+                      {i18n.translate('xpack.observabilityAiAssistant.prompt.emptySelection', {
+                        defaultMessage: 'Empty selection',
+                      })}
+                    </EuiButtonEmpty>
+                  ) : null}
+                </EuiFlexItem>
+              </EuiFlexGroup>
+            </EuiFlexItem>
+            <EuiFlexItem>
+              {selectedFunctionName ? (
+                <EuiPanel borderRadius="none" color="subdued" hasShadow={false} paddingSize="xs">
+                  <CodeEditor
+                    aria-label="payloadEditor"
+                    fullWidth
+                    height="120px"
+                    languageId="json"
+                    isCopyable
+                    languageConfiguration={{
+                      autoClosingPairs: [
+                        {
+                          open: '{',
+                          close: '}',
+                        },
+                      ],
+                    }}
+                    editorDidMount={(editor) => {
+                      editor.focus();
+                    }}
+                    options={{
+                      accessibilitySupport: 'off',
+                      acceptSuggestionOnEnter: 'on',
+                      automaticLayout: true,
+                      autoClosingQuotes: 'always',
+                      autoIndent: 'full',
+                      contextmenu: true,
+                      fontSize: 12,
+                      formatOnPaste: true,
+                      formatOnType: true,
+                      inlineHints: { enabled: true },
+                      lineNumbers: 'on',
+                      minimap: { enabled: false },
+                      model,
+                      overviewRulerBorder: false,
+                      quickSuggestions: true,
+                      scrollbar: { alwaysConsumeMouseWheel: false },
+                      scrollBeyondLastLine: false,
+                      suggestOnTriggerCharacters: true,
+                      tabSize: 2,
+                      wordWrap: 'on',
+                      wrappingIndent: 'indent',
+                    }}
+                    transparentBackground
+                    value={functionPayload || ''}
+                    onChange={handleChangeFunctionPayload}
+                  />
+                </EuiPanel>
+              ) : (
+                <EuiTextArea
                   fullWidth
-                  height="120px"
-                  languageId="json"
-                  value={functionPayload || ''}
-                  onChange={handleChangeFunctionPayload}
-                  isCopyable
-                  languageConfiguration={{
-                    autoClosingPairs: [
-                      {
-                        open: '{',
-                        close: '}',
-                      },
-                    ],
-                  }}
-                  options={{
-                    accessibilitySupport: 'off',
-                    acceptSuggestionOnEnter: 'on',
-                    automaticLayout: true,
-                    autoClosingQuotes: 'always',
-                    autoIndent: 'full',
-                    contextmenu: true,
-                    fontSize: 12,
-                    formatOnPaste: true,
-                    formatOnType: true,
-                    inlineHints: { enabled: true },
-                    lineNumbers: 'on',
-                    minimap: { enabled: false },
-                    model,
-                    overviewRulerBorder: false,
-                    quickSuggestions: true,
-                    scrollbar: { alwaysConsumeMouseWheel: false },
-                    scrollBeyondLastLine: false,
-                    suggestOnTriggerCharacters: true,
-                    tabSize: 2,
-                    wordWrap: 'on',
-                    wrappingIndent: 'indent',
-                  }}
-                  transparentBackground
+                  inputRef={textAreaRef}
+                  placeholder={i18n.translate('xpack.observabilityAiAssistant.prompt.placeholder', {
+                    defaultMessage: 'Press ‘$’ for function recommendations',
+                  })}
+                  resize="vertical"
+                  rows={1}
+                  value={prompt}
+                  onChange={handleChange}
                 />
-              </EuiPanel>
-            ) : (
-              <EuiFieldText
-                fullWidth
-                value={prompt}
-                placeholder={i18n.translate('xpack.observabilityAiAssistant.prompt.placeholder', {
-                  defaultMessage: 'Press ‘$’ for function recommendations',
-                })}
-                inputRef={ref}
-                onChange={handleChange}
-                onSubmit={handleSubmit}
-              />
-            )}
-          </EuiFlexItem>
-        </EuiFlexGroup>
-      </EuiFlexItem>
-      <EuiFlexItem grow={false}>
-        <EuiSpacer size="xl" />
-        <EuiButtonIcon
-          aria-label="Submit"
-          isLoading={loading}
-          disabled={selectedFunction ? false : !prompt || loading || disabled}
-          display={
-            selectedFunction ? (functionPayload ? 'fill' : 'base') : prompt ? 'fill' : 'base'
-          }
-          iconType="kqlFunction"
-          size="m"
-          onClick={handleSubmit}
-        />
-      </EuiFlexItem>
-    </EuiFlexGroup>
+              )}
+            </EuiFlexItem>
+          </EuiFlexGroup>
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiSpacer size="xl" />
+          <EuiButtonIcon
+            aria-label="Submit"
+            disabled={selectedFunctionName ? false : !prompt || loading || disabled}
+            display={
+              selectedFunctionName ? (functionPayload ? 'fill' : 'base') : prompt ? 'fill' : 'base'
+            }
+            iconType="kqlFunction"
+            isLoading={loading}
+            size="m"
+            onClick={handleSubmit}
+          />
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    </EuiFocusTrap>
   );
 }

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.stories.tsx

@@ -77,12 +77,16 @@ const defaultProps: ComponentProps<typeof Component> = {
         arguments: '{ "foo": "bar" }',
         trigger: MessageRole.Assistant,
       },
-      canEdit: true,
+      actions: {
+        canEdit: true,
+      },
     }),
     buildFunctionChatItem({
       content: '{ "message": "The arguments are wrong" }',
       error: new Error(),
-      canRegenerate: false,
+      actions: {
+        canRegenerate: false,
+      },
     }),
     buildAssistantChatItem({
       content: '',
@@ -92,7 +96,9 @@ const defaultProps: ComponentProps<typeof Component> = {
         arguments: '{ "bar": "foo" }',
         trigger: MessageRole.Assistant,
       },
-      canEdit: true,
+      actions: {
+        canEdit: true,
+      },
     }),
     buildFunctionChatItem({
       content: '',
@@ -100,7 +106,7 @@ const defaultProps: ComponentProps<typeof Component> = {
       loading: true,
     }),
   ],
-  onEdit: () => {},
+  onEdit: async () => {},
   onFeedback: () => {},
   onRegenerate: () => {},
   onStopGenerating: () => {},

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.tsx

@@ -7,26 +7,36 @@
 
 import { EuiCommentList } from '@elastic/eui';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
-import React from 'react';
-import type { Message } from '../../../common';
+import { compact } from 'lodash';
+import React, { ReactNode } from 'react';
+import { type Message } from '../../../common';
+
 import type { Feedback } from '../feedback_buttons';
 import { ChatItem } from './chat_item';
 
 export interface ChatTimelineItem
   extends Pick<Message['message'], 'role' | 'content' | 'function_call'> {
   id: string;
-  title: string;
+  title: ReactNode;
+  actions: {
+    canCopy: boolean;
+    canEdit: boolean;
+    canGiveFeedback: boolean;
+    canRegenerate: boolean;
+  };
+  display: {
+    collapsed: boolean;
+    hide?: boolean;
+  };
   loading: boolean;
-  error?: any;
-  canEdit: boolean;
-  canRegenerate: boolean;
-  canGiveFeedback: boolean;
+  element?: React.ReactNode;
   currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
+  error?: any;
 }
 
 export interface ChatTimelineProps {
   items: ChatTimelineItem[];
-  onEdit: (item: ChatTimelineItem, content: string) => void;
+  onEdit: (item: ChatTimelineItem, message: Message) => Promise<void>;
   onFeedback: (item: ChatTimelineItem, feedback: Feedback) => void;
   onRegenerate: (item: ChatTimelineItem) => void;
   onStopGenerating: () => void;
@@ -41,23 +51,27 @@ export function ChatTimeline({
 }: ChatTimelineProps) {
   return (
     <EuiCommentList>
-      {items.map((item, index) => (
-        <ChatItem
-          // use index, not id to prevent unmounting of component when message is persisted
-          key={index}
-          {...item}
-          onFeedbackClick={(feedback) => {
-            onFeedback(item, feedback);
-          }}
-          onRegenerateClick={() => {
-            onRegenerate(item);
-          }}
-          onEditSubmit={(content) => {
-            onEdit(item, content);
-          }}
-          onStopGeneratingClick={onStopGenerating}
-        />
-      ))}
+      {compact(
+        items.map((item, index) =>
+          !item.display.hide ? (
+            <ChatItem
+              // use index, not id to prevent unmounting of component when message is persisted
+              key={index}
+              {...item}
+              onFeedbackClick={(feedback) => {
+                onFeedback(item, feedback);
+              }}
+              onRegenerateClick={() => {
+                onRegenerate(item);
+              }}
+              onEditSubmit={(message) => {
+                return onEdit(item, message);
+              }}
+              onStopGeneratingClick={onStopGenerating}
+            />
+          ) : null
+        )
+      )}
     </EuiCommentList>
   );
 }

x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx

@@ -37,11 +37,11 @@ export function ConversationList({
   onClickDeleteConversation,
 }: {
   selected: string;
-  onClickConversation: (conversationId: string) => void;
-  onClickNewChat: () => void;
   loading: boolean;
   error?: any;
   conversations?: Array<{ id: string; label: string; href?: string }>;
+  onClickConversation: (conversationId: string) => void;
+  onClickNewChat: () => void;
   onClickDeleteConversation: (id: string) => void;
 }) {
   return (

x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.stories.tsx

@@ -23,7 +23,6 @@ const Template: ComponentStory<typeof Component> = (props: FunctionListPopover)
 };
 
 const defaultProps: FunctionListPopover = {
-  functions: [],
   onSelectFunction: () => {},
 };
 

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

@@ -8,7 +8,7 @@
 import React, { useEffect, useState } from 'react';
 import {
   EuiButtonEmpty,
-  EuiContextMenuItem,
+  EuiContextMenu,
   EuiContextMenuPanel,
   EuiPopover,
   EuiSpacer,
@@ -16,16 +16,17 @@ import {
 } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 import { FunctionDefinition } from '../../../common/types';
+import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
 
 export function FunctionListPopover({
-  functions,
-  selectedFunction,
+  selectedFunctionName,
   onSelectFunction,
 }: {
-  functions: FunctionDefinition[];
-  selectedFunction?: FunctionDefinition;
-  onSelectFunction: (func: FunctionDefinition) => void;
+  selectedFunctionName?: string;
+  onSelectFunction: (func: string) => void;
 }) {
+  const chatService = useObservabilityAIAssistantChatService();
+
   const [isFunctionListOpen, setIsFunctionListOpen] = useState(false);
 
   const handleClickFunctionList = () => {
@@ -34,7 +35,7 @@ export function FunctionListPopover({
 
   const handleSelectFunction = (func: FunctionDefinition) => {
     setIsFunctionListOpen(false);
-    onSelectFunction(func);
+    onSelectFunction(func.options.name);
   };
 
   useEffect(() => {
@@ -61,31 +62,44 @@ export function FunctionListPopover({
           size="xs"
           onClick={handleClickFunctionList}
         >
-          {selectedFunction
-            ? selectedFunction.options.name
+          {selectedFunctionName
+            ? selectedFunctionName
             : i18n.translate('xpack.observabilityAiAssistant.prompt.callFunction', {
                 defaultMessage: 'Call function',
               })}
         </EuiButtonEmpty>
       }
       closePopover={handleClickFunctionList}
+      css={{ maxWidth: 400 }}
       panelPaddingSize="none"
       isOpen={isFunctionListOpen}
     >
       <EuiContextMenuPanel size="s">
-        {functions.map((func) => (
-          <EuiContextMenuItem key={func.options.name} onClick={() => handleSelectFunction(func)}>
-            <EuiText size="s">
-              <p>
-                <strong>{func.options.name}</strong>
-              </p>
-            </EuiText>
-            <EuiSpacer size="xs" />
-            <EuiText size="s">
-              <p>{func.options.description}</p>
-            </EuiText>
-          </EuiContextMenuItem>
-        ))}
+        <EuiContextMenu
+          initialPanelId={0}
+          panels={[
+            {
+              id: 0,
+              width: 500,
+              items: chatService.getFunctions().map((func) => ({
+                name: (
+                  <>
+                    <EuiText size="s">
+                      <p>
+                        <strong>{func.options.name}</strong>
+                      </p>
+                    </EuiText>
+                    <EuiSpacer size="xs" />
+                    <EuiText size="s">
+                      <p>{func.options.descriptionForUser}</p>
+                    </EuiText>
+                  </>
+                ),
+                onClick: () => handleSelectFunction(func),
+              })),
+            },
+          ]}
+        />
       </EuiContextMenuPanel>
     </EuiPopover>
   );

x-pack/plugins/observability_ai_assistant/public/components/chat/knowledge_base_callout.stories.tsx

@@ -0,0 +1,70 @@
+/*
+ * 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 { ComponentMeta, ComponentStoryObj } from '@storybook/react';
+import { merge } from 'lodash';
+import { KibanaReactStorybookDecorator } from '../../utils/storybook_decorator';
+import { KnowledgeBaseCallout as Component } from './knowledge_base_callout';
+
+const meta: ComponentMeta<typeof Component> = {
+  component: Component,
+  title: 'app/Molecules/KnowledgeBaseCallout',
+  decorators: [KibanaReactStorybookDecorator],
+};
+
+export default meta;
+const defaultProps: ComponentStoryObj<typeof Component> = {
+  args: {
+    knowledgeBase: {
+      status: {
+        loading: false,
+        value: {
+          ready: false,
+        },
+        refresh: () => {},
+      },
+      isInstalling: false,
+      installError: undefined,
+      install: async () => {},
+    },
+  },
+};
+
+export const StatusError: ComponentStoryObj<typeof Component> = merge({}, defaultProps, {
+  args: { knowledgeBase: { status: { loading: false, error: new Error() } } },
+});
+
+export const Loading: ComponentStoryObj<typeof Component> = merge({}, defaultProps, {
+  args: { knowledgeBase: { status: { loading: true } } },
+});
+
+export const NotInstalled: ComponentStoryObj<typeof Component> = merge({}, defaultProps, {
+  args: { knowledgeBase: { status: { loading: false, value: { ready: false } } } },
+});
+
+export const Installing: ComponentStoryObj<typeof Component> = merge({}, defaultProps, {
+  args: {
+    knowledgeBase: { status: { loading: false, value: { ready: false } }, isInstalling: true },
+  },
+});
+
+export const InstallError: ComponentStoryObj<typeof Component> = merge({}, defaultProps, {
+  args: {
+    knowledgeBase: {
+      status: {
+        loading: false,
+        value: { ready: false },
+      },
+      isInstalling: false,
+      installError: new Error(),
+    },
+  },
+});
+
+export const Installed: ComponentStoryObj<typeof Component> = merge({}, defaultProps, {
+  args: { knowledgeBase: { status: { loading: false, value: { ready: true } } } },
+});

x-pack/plugins/observability_ai_assistant/public/components/chat/knowledge_base_callout.tsx

@@ -0,0 +1,107 @@
+/*
+ * 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 React from 'react';
+import {
+  EuiFlexGroup,
+  EuiFlexItem,
+  EuiLink,
+  EuiLoadingSpinner,
+  EuiPanel,
+  EuiText,
+} from '@elastic/eui';
+import { i18n } from '@kbn/i18n';
+import { UseKnowledgeBaseResult } from '../../hooks/use_knowledge_base';
+
+export function KnowledgeBaseCallout({ knowledgeBase }: { knowledgeBase: UseKnowledgeBaseResult }) {
+  let content: React.ReactNode;
+
+  let color: 'primary' | 'danger' | 'plain' = 'primary';
+
+  if (knowledgeBase.status.loading) {
+    content = (
+      <EuiFlexGroup alignItems="center" gutterSize="s">
+        <EuiFlexItem grow={false}>
+          <EuiLoadingSpinner size="s" />
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiText size="xs" color="subdued">
+            {i18n.translate('xpack.observabilityAiAssistant.checkingKbAvailability', {
+              defaultMessage: 'Checking availability of knowledge base',
+            })}
+          </EuiText>
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    );
+  } else if (knowledgeBase.status.error) {
+    color = 'danger';
+    content = (
+      <EuiText size="xs" color={color}>
+        {i18n.translate('xpack.observabilityAiAssistant.failedToGetStatus', {
+          defaultMessage: 'Failed to get model status.',
+        })}
+      </EuiText>
+    );
+  } else if (knowledgeBase.status.value?.ready) {
+    color = 'plain';
+    content = (
+      <EuiText size="xs" color="subdued">
+        {i18n.translate('xpack.observabilityAiAssistant.poweredByModel', {
+          defaultMessage: 'Powered by {model}',
+          values: {
+            model: 'ELSER',
+          },
+        })}
+      </EuiText>
+    );
+  } else if (knowledgeBase.isInstalling) {
+    color = 'primary';
+    content = (
+      <EuiFlexGroup alignItems="center" gutterSize="s">
+        <EuiFlexItem grow={false}>
+          <EuiLoadingSpinner size="s" />
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiText size="xs" color={color}>
+            {i18n.translate('xpack.observabilityAiAssistant.installingKb', {
+              defaultMessage: 'Setting up the knowledge base',
+            })}
+          </EuiText>
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    );
+  } else if (knowledgeBase.installError) {
+    color = 'danger';
+    content = (
+      <EuiText size="xs" color={color}>
+        {i18n.translate('xpack.observabilityAiAssistant.failedToSetupKnowledgeBase', {
+          defaultMessage: 'Failed to set up knowledge base.',
+        })}
+      </EuiText>
+    );
+  } else if (!knowledgeBase.status.value?.ready && !knowledgeBase.status.error) {
+    content = (
+      <EuiLink
+        onClick={() => {
+          knowledgeBase.install();
+        }}
+      >
+        <EuiText size="xs">
+          {i18n.translate('xpack.observabilityAiAssistant.setupKb', {
+            defaultMessage: 'Improve your experience by setting up the knowledge base.',
+          })}
+        </EuiText>
+      </EuiLink>
+    );
+  }
+
+  return (
+    <EuiPanel hasBorder={false} hasShadow={false} borderRadius="none" color={color} paddingSize="s">
+      {content}
+    </EuiPanel>
+  );
+}

x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx

@@ -11,7 +11,6 @@ import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import type { Subscription } from 'rxjs';
 import { MessageRole, type Message } from '../../../common/types';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
-import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import type { PendingMessage } from '../../types';
 import { ChatFlyout } from '../chat/chat_flyout';
 import { ConnectorSelectorBase } from '../connector_selector/connector_selector_base';
@@ -23,6 +22,10 @@ import { StopGeneratingButton } from '../buttons/stop_generating_button';
 import { InsightBase } from './insight_base';
 import { MissingCredentialsCallout } from '../missing_credentials_callout';
 import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
+import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
+import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
+import { useAbortableAsync } from '../../hooks/use_abortable_async';
+import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
 
 function ChatContent({
   title,
@@ -33,7 +36,7 @@ function ChatContent({
   messages: Message[];
   connectorId: string;
 }) {
-  const service = useObservabilityAIAssistant();
+  const chatService = useObservabilityAIAssistantChatService();
 
   const [pendingMessage, setPendingMessage] = useState<PendingMessage | undefined>();
   const [loading, setLoading] = useState(false);
@@ -42,7 +45,7 @@ function ChatContent({
   const reloadReply = useCallback(() => {
     setLoading(true);
 
-    const nextSubscription = service.chat({ messages, connectorId }).subscribe({
+    const nextSubscription = chatService.chat({ messages, connectorId }).subscribe({
       next: (msg) => {
         setPendingMessage(() => msg);
       },
@@ -52,7 +55,7 @@ function ChatContent({
     });
 
     setSubscription(nextSubscription);
-  }, [messages, connectorId, service]);
+  }, [messages, connectorId, chatService]);
 
   useEffect(() => {
     reloadReply();
@@ -129,6 +132,15 @@ export function Insight({ messages, title }: { messages: Message[]; title: strin
 
   const connectors = useGenAIConnectors();
 
+  const service = useObservabilityAIAssistant();
+
+  const chatService = useAbortableAsync(
+    ({ signal }) => {
+      return service.start({ signal });
+    },
+    [service]
+  );
+
   const {
     services: { http },
   } = useKibana();
@@ -152,9 +164,13 @@ export function Insight({ messages, title }: { messages: Message[]; title: strin
         setHasOpened((prevHasOpened) => prevHasOpened || isOpen);
       }}
       controls={<ConnectorSelectorBase {...connectors} />}
-      loading={connectors.loading}
+      loading={connectors.loading || chatService.loading}
     >
-      {children}
+      {chatService.value ? (
+        <ObservabilityAIAssistantChatServiceProvider value={chatService.value}>
+          {children}
+        </ObservabilityAIAssistantChatServiceProvider>
+      ) : null}
     </InsightBase>
   );
 }

x-pack/plugins/observability_ai_assistant/public/components/message_panel/message_text.tsx

@@ -87,7 +87,7 @@ export function MessageText(props: Props) {
   const containerClassName = css`
     overflow-wrap: break-word;
 
-    code {
+    pre {
       background: ${euiThemeVars.euiColorLightestShade};
       padding: 0 8px;
     }

x-pack/plugins/observability_ai_assistant/public/components/page_template.tsx

@@ -14,6 +14,7 @@ const pageSectionContentClassName = css`
   flex-grow: 1;
   padding-top: 0;
   padding-bottom: 0;
+  max-block-size: calc(100vh - 96px);
 `;
 
 export function ObservabilityAIAssistantPageTemplate({ children }: { children: React.ReactNode }) {

x-pack/plugins/observability_ai_assistant/public/components/render_function.tsx

@@ -0,0 +1,21 @@
+/*
+ * 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 React from 'react';
+import { Message } from '../../common';
+import { useObservabilityAIAssistantChatService } from '../hooks/use_observability_ai_assistant_chat_service';
+
+interface Props {
+  name: string;
+  arguments: string | undefined;
+  response: Message['message'];
+}
+
+export function RenderFunction(props: Props) {
+  const chatService = useObservabilityAIAssistantChatService();
+
+  return <>{chatService.renderFunction(props.name, props.arguments, props.response)}</>;
+}

x-pack/plugins/observability_ai_assistant/public/context/observability_ai_assistant_chat_service_provider.tsx

@@ -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 { createContext } from 'react';
+import type { ObservabilityAIAssistantChatService } from '../types';
+
+export const ObservabilityAIAssistantChatServiceContext = createContext<
+  ObservabilityAIAssistantChatService | undefined
+>(undefined);
+
+export const ObservabilityAIAssistantChatServiceProvider =
+  ObservabilityAIAssistantChatServiceContext.Provider;

x-pack/plugins/observability_ai_assistant/public/functions/elasticsearch.ts

@@ -21,6 +21,7 @@ export function registerElasticsearchFunction({
       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: {
@@ -34,7 +35,7 @@ export function registerElasticsearchFunction({
             description: 'The path of the Elasticsearch endpoint, including query parameters',
           },
         },
-        required: ['method' as const, 'path' as const],
+        required: ['method', 'path'] as const,
       },
     },
     ({ arguments: { method, path, body } }, signal) => {

x-pack/plugins/observability_ai_assistant/public/functions/index.ts

@@ -5,30 +5,68 @@
  * 2.0.
  */
 
+import dedent from 'dedent';
 import type { RegisterContextDefinition, RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantPluginStartDependencies } from '../types';
 import type { ObservabilityAIAssistantService } from '../types';
 import { registerElasticsearchFunction } from './elasticsearch';
 import { registerRecallFunction } from './recall';
-import { registerSetupKbFunction } from './setup_kb';
 import { registerSummarisationFunction } from './summarise';
 
-export function registerFunctions({
+export async function registerFunctions({
   registerFunction,
   registerContext,
   service,
+  pluginsStart,
+  signal,
 }: {
   registerFunction: RegisterFunctionDefinition;
   registerContext: RegisterContextDefinition;
   service: ObservabilityAIAssistantService;
+  pluginsStart: ObservabilityAIAssistantPluginStartDependencies;
+  signal: AbortSignal;
 }) {
-  registerContext({
-    name: 'core',
-    description:
-      'Core functions, like calling Elasticsearch APIs, storing embeddables for instructions or creating base visualisations.',
-  });
+  return service
+    .callApi('GET /internal/observability_ai_assistant/functions/kb_status', {
+      signal,
+    })
+    .then((response) => {
+      const isReady = response.ready;
 
-  registerElasticsearchFunction({ service, registerFunction });
-  registerSummarisationFunction({ service, registerFunction });
-  registerRecallFunction({ service, registerFunction });
-  registerSetupKbFunction({ service, registerFunction });
+      let description = `You have the ability to call Elasticsearch APIs or create visualisations using Lens.`;
+
+      if (isReady) {
+        description += `You can use the "summarise" 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. When you create this summarisation, make sure you craft it in a way that can be recalled with a semantic search later.
+
+        Additionally, you can use the "recall" function to retrieve relevant information from the knowledge database. Using the "recall" function will allow you to create a much more tailored and improved user experience. Assume the user has told you before.
+      
+        For every question that the user asks, you, as the assistant, must query the knowledge base, no matter how confident you are in your response.
+      
+        When the user starts a conversation, you as the assistant, MUST ALWAYS use the "recall" function to make sure you have enough context from the knowledge database to understand their question. DO NOT assume your knowledge is sufficient to answer the question appropriately. Make sure you mention this to the user in your first reply to them.
+      
+        You do have access to private data that the user has shared with you in previous conversations.
+      
+        Given this ability to store and recall data, it is very important to not make assumptions about the user's data or preferences, but rather, first query the knowledge database, and if nothing relevant comes up, ask the user for clarification. This is very important! They could be storing their data in any field, and in any data stream or index.
+      
+        RIGHT:
+        User: "What is NASA"
+        Assistant executes recall function
+        Assistant answers question with data from recall function response
+      
+        WRONG:
+        User: "What is NASA"
+        Assistant answers question without querying the knowledge`;
+        registerSummarisationFunction({ service, registerFunction });
+        registerRecallFunction({ service, registerFunction });
+      } else {
+        description += `You do not have a working memory. Don't try to recall information via other functions.  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 });
+
+      registerContext({
+        name: 'core',
+        description: dedent(description),
+      });
+    });
 }

x-pack/plugins/observability_ai_assistant/public/functions/recall.ts

@@ -22,6 +22,7 @@ export function registerRecallFunction({
       contexts: ['core'],
       description:
         'Use this function to recall earlier learnings. Anything you will summarise can be retrieved again later via this function.',
+      descriptionForUser: 'This function allows the assistant to recall previous learnings.',
       parameters: {
         type: 'object',
         properties: {

x-pack/plugins/observability_ai_assistant/public/functions/setup_kb.ts

@@ -1,38 +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 { Serializable } from '@kbn/utility-types';
-import type { RegisterFunctionDefinition } from '../../common/types';
-import type { ObservabilityAIAssistantService } from '../types';
-
-export function registerSetupKbFunction({
-  service,
-  registerFunction,
-}: {
-  service: ObservabilityAIAssistantService;
-  registerFunction: RegisterFunctionDefinition;
-}) {
-  registerFunction(
-    {
-      name: 'setup_kb',
-      contexts: ['core'],
-      description:
-        'Use this function to set up the knowledge base. ONLY use this if you got an error from the recall or summarise function, or if the user has explicitly requested it. Note that it might take a while (e.g. ten minutes) until the knowledge base is available. Assume it will not be ready for the rest of the current conversation.',
-      parameters: {
-        type: 'object',
-        properties: {},
-      },
-    },
-    ({}, signal) => {
-      return service
-        .callApi('POST /internal/observability_ai_assistant/functions/setup_kb', {
-          signal,
-        })
-        .then((response) => ({ content: response as unknown as Serializable }));
-    }
-  );
-}

x-pack/plugins/observability_ai_assistant/public/functions/summarise.ts

@@ -21,6 +21,8 @@ export function registerSummarisationFunction({
       contexts: ['core'],
       description:
         'Use this function to summarise things learned from the conversation. You can score the learnings with a confidence metric, whether it is a correction on a previous learning. An embedding will be created that you can recall later with a semantic search. There is no need to ask the user for permission to store something you have learned, unless you do not feel confident.',
+      descriptionForUser:
+        'This function allows the Elastic Assistant to summarise things from the conversation.',
       parameters: {
         type: 'object',
         properties: {

x-pack/plugins/observability_ai_assistant/public/hooks/__storybook_mocks__/use_knowledge_base.ts

@@ -0,0 +1,23 @@
+/*
+ * 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 { UseKnowledgeBaseResult } from '../use_knowledge_base';
+
+export function useKnowledgeBase(): UseKnowledgeBaseResult {
+  return {
+    install: async () => {},
+    isInstalling: false,
+    status: {
+      loading: false,
+      refresh: () => {},
+      error: undefined,
+      value: {
+        ready: true,
+      },
+    },
+  };
+}

x-pack/plugins/observability_ai_assistant/public/hooks/__storybook_mocks__/use_observability_ai_assistant.ts

@@ -5,6 +5,14 @@
  * 2.0.
  */
 
+const service = {
+  start: async () => {
+    return {
+      getFunctions: [],
+    };
+  },
+};
+
 export function useObservabilityAIAssistant() {
-  return {};
+  return service;
 }

x-pack/plugins/observability_ai_assistant/public/hooks/__storybook_mocks__/use_observability_ai_assistant_chat_service.ts

@@ -0,0 +1,14 @@
+/*
+ * 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.
+ */
+
+export function useObservabilityAIAssistantChatService() {
+  return {
+    getFunctions: () => {
+      return [];
+    },
+  };
+}

x-pack/plugins/observability_ai_assistant/public/hooks/use_abortable_async.ts

@@ -49,11 +49,11 @@ export function useAbortableAsync<T>(
       } else {
         setError(undefined);
         setValue(response);
+        setLoading(false);
       }
     } catch (err) {
       setValue(undefined);
       setError(err);
-    } finally {
       setLoading(false);
     }
 

x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts

@@ -9,12 +9,19 @@ import { merge, omit } from 'lodash';
 import { Dispatch, SetStateAction, useState } from 'react';
 import type { Conversation, Message } from '../../common';
 import type { ConversationCreateRequest } from '../../common/types';
+import { ObservabilityAIAssistantChatService } from '../types';
 import { useAbortableAsync, type AbortableAsyncState } from './use_abortable_async';
 import { useKibana } from './use_kibana';
 import { useObservabilityAIAssistant } from './use_observability_ai_assistant';
 import { createNewConversation } from './use_timeline';
 
-export function useConversation(conversationId?: string): {
+export function useConversation({
+  conversationId,
+  chatService,
+}: {
+  conversationId?: string;
+  chatService?: ObservabilityAIAssistantChatService;
+}): {
   conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined>;
   displayedMessages: Message[];
   setDisplayedMessages: Dispatch<SetStateAction<Message[]>>;
@@ -32,7 +39,9 @@ export function useConversation(conversationId?: string): {
     useAbortableAsync(
       ({ signal }) => {
         if (!conversationId) {
-          const nextConversation = createNewConversation();
+          const nextConversation = createNewConversation({
+            contexts: chatService?.getContexts() || [],
+          });
           setDisplayedMessages(nextConversation.messages);
           return nextConversation;
         }
@@ -51,7 +60,7 @@ export function useConversation(conversationId?: string): {
             throw error;
           });
       },
-      [conversationId]
+      [conversationId, chatService]
     );
 
   return {

x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts

@@ -4,16 +4,29 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import { useMemo } from 'react';
 import { monaco } from '@kbn/monaco';
-import { FunctionDefinition } from '../../common/types';
+import { useMemo } from 'react';
+import { createInitializedObject } from '../utils/create_initialized_object';
+import { useObservabilityAIAssistantChatService } from './use_observability_ai_assistant_chat_service';
 
 const { editor, languages, Uri } = monaco;
 
 const SCHEMA_URI = 'http://elastic.co/foo.json';
 const modelUri = Uri.parse(SCHEMA_URI);
 
-export const useJsonEditorModel = (functionDefinition?: FunctionDefinition) => {
+export const useJsonEditorModel = ({
+  functionName,
+  initialJson,
+}: {
+  functionName: string | undefined;
+  initialJson?: string | undefined;
+}) => {
+  const chatService = useObservabilityAIAssistantChatService();
+
+  const functionDefinition = chatService
+    .getFunctions()
+    .find((func) => func.options.name === functionName);
+
   return useMemo(() => {
     if (!functionDefinition) {
       return {};
@@ -21,14 +34,10 @@ export const useJsonEditorModel = (functionDefinition?: FunctionDefinition) => {
 
     const schema = { ...functionDefinition.options.parameters };
 
-    const initialJsonString = functionDefinition.options.parameters.properties
-      ? Object.keys(functionDefinition.options.parameters.properties).reduce(
-          (acc, curr, index, arr) => {
-            const val = `${acc}  "${curr}": "",\n`;
-            return index === arr.length - 1 ? `${val}}` : val;
-          },
-          '{\n'
-        )
+    const initialJsonString = initialJson
+      ? initialJson
+      : functionDefinition.options.parameters.properties
+      ? JSON.stringify(createInitializedObject(functionDefinition.options.parameters), null, 4)
       : '';
 
     languages.json.jsonDefaults.setDiagnosticsOptions({
@@ -49,5 +58,5 @@ export const useJsonEditorModel = (functionDefinition?: FunctionDefinition) => {
     }
 
     return { model, initialJsonString };
-  }, [functionDefinition]);
+  }, [functionDefinition, initialJson]);
 };

x-pack/plugins/observability_ai_assistant/public/hooks/use_knowledge_base.ts

@@ -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 { i18n } from '@kbn/i18n';
+import { useMemo, useState } from 'react';
+import { AbortableAsyncState, useAbortableAsync } from './use_abortable_async';
+import { useKibana } from './use_kibana';
+import { useObservabilityAIAssistant } from './use_observability_ai_assistant';
+
+export interface UseKnowledgeBaseResult {
+  status: AbortableAsyncState<{
+    ready: boolean;
+    error?: any;
+    deployment_state?: string;
+    allocation_state?: string;
+  }>;
+  isInstalling: boolean;
+  installError?: Error;
+  install: () => Promise<void>;
+}
+
+export function useKnowledgeBase(): UseKnowledgeBaseResult {
+  const {
+    notifications: { toasts },
+  } = useKibana().services;
+  const service = useObservabilityAIAssistant();
+
+  const status = useAbortableAsync(({ signal }) => {
+    return service.callApi('GET /internal/observability_ai_assistant/functions/kb_status', {
+      signal,
+    });
+  }, []);
+
+  const [isInstalling, setIsInstalling] = useState(false);
+
+  const [installError, setInstallError] = useState<Error>();
+
+  return useMemo(
+    () => ({
+      status,
+      isInstalling,
+      installError,
+      install: () => {
+        setIsInstalling(true);
+        return service
+          .callApi('POST /internal/observability_ai_assistant/functions/setup_kb', {
+            signal: null,
+          })
+          .then(() => {
+            status.refresh();
+          })
+          .catch((error) => {
+            setInstallError(error);
+            toasts.addError(error, {
+              title: i18n.translate('xpack.observabilityAiAssistant.errorSettingUpKnowledgeBase', {
+                defaultMessage: 'Could not set up Knowledge Base',
+              }),
+            });
+          })
+          .finally(() => {
+            setIsInstalling(false);
+          });
+      },
+    }),
+    [status, isInstalling, installError, service, toasts]
+  );
+}

x-pack/plugins/observability_ai_assistant/public/hooks/use_observability_ai_assistant_chat_service.ts

@@ -0,0 +1,20 @@
+/*
+ * 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 { useContext } from 'react';
+import { ObservabilityAIAssistantChatServiceContext } from '../context/observability_ai_assistant_chat_service_provider';
+
+export function useObservabilityAIAssistantChatService() {
+  const services = useContext(ObservabilityAIAssistantChatServiceContext);
+
+  if (!services) {
+    throw new Error(
+      'ObservabilityAIAssistantChatServiceContext not set. Did you wrap your component in `<ObservabilityAIAssistantChatServiceProvider/>`?'
+    );
+  }
+
+  return services;
+}

x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.test.ts

@@ -34,7 +34,7 @@ describe('useTimeline', () => {
             selectConnector: () => {},
             connectors: [{ id: 'OpenAI' }] as FindActionResult[],
           },
-          service: {},
+          chatService: {},
           messages: [],
           onChatComplete: jest.fn(),
           onChatUpdate: jest.fn(),
@@ -45,9 +45,16 @@ describe('useTimeline', () => {
       expect(hookResult.result.current.items.length).toEqual(1);
 
       expect(hookResult.result.current.items[0]).toEqual({
-        canEdit: false,
-        canRegenerate: false,
-        canGiveFeedback: false,
+        display: {
+          collapsed: false,
+          hide: false,
+        },
+        actions: {
+          canCopy: false,
+          canEdit: false,
+          canRegenerate: false,
+          canGiveFeedback: false,
+        },
         role: MessageRole.User,
         title: 'started a conversation',
         loading: false,
@@ -62,14 +69,12 @@ describe('useTimeline', () => {
         initialProps: {
           messages: [
             {
-              '@timestamp': new Date().toISOString(),
               message: {
                 role: MessageRole.User,
                 content: 'Hello',
               },
             },
             {
-              '@timestamp': new Date().toISOString(),
               message: {
                 role: MessageRole.Assistant,
                 content: 'Goodbye',
@@ -79,7 +84,7 @@ describe('useTimeline', () => {
           connectors: {
             selectedConnector: 'foo',
           },
-          service: {
+          chatService: {
             chat: () => {},
           },
         } as unknown as HookProps,
@@ -89,9 +94,16 @@ describe('useTimeline', () => {
       expect(hookResult.result.current.items.length).toEqual(3);
 
       expect(hookResult.result.current.items[1]).toEqual({
-        canEdit: true,
-        canRegenerate: false,
-        canGiveFeedback: false,
+        actions: {
+          canCopy: true,
+          canEdit: true,
+          canRegenerate: false,
+          canGiveFeedback: false,
+        },
+        display: {
+          collapsed: false,
+          hide: false,
+        },
         role: MessageRole.User,
         content: 'Hello',
         loading: false,
@@ -100,9 +112,16 @@ describe('useTimeline', () => {
       });
 
       expect(hookResult.result.current.items[2]).toEqual({
-        canEdit: false,
-        canRegenerate: true,
-        canGiveFeedback: true,
+        display: {
+          collapsed: false,
+          hide: false,
+        },
+        actions: {
+          canCopy: true,
+          canEdit: false,
+          canRegenerate: true,
+          canGiveFeedback: true,
+        },
         role: MessageRole.Assistant,
         content: 'Goodbye',
         loading: false,
@@ -115,11 +134,11 @@ describe('useTimeline', () => {
   describe('when submitting a new prompt', () => {
     let subject: Subject<PendingMessage>;
 
-    let props: Omit<HookProps, 'onChatUpdate' | 'onChatComplete' | 'service'> & {
+    let props: Omit<HookProps, 'onChatUpdate' | 'onChatComplete' | 'chatService'> & {
       onChatUpdate: jest.MockedFn<HookProps['onChatUpdate']>;
       onChatComplete: jest.MockedFn<HookProps['onChatComplete']>;
-      service: Omit<HookProps['service'], 'executeFunction'> & {
-        executeFunction: jest.MockedFn<HookProps['service']['executeFunction']>;
+      chatService: Omit<HookProps['chatService'], 'executeFunction'> & {
+        executeFunction: jest.MockedFn<HookProps['chatService']['executeFunction']>;
       };
     };
 
@@ -129,7 +148,7 @@ describe('useTimeline', () => {
         connectors: {
           selectedConnector: 'foo',
         },
-        service: {
+        chatService: {
           chat: jest.fn().mockImplementation(() => {
             subject = new BehaviorSubject<PendingMessage>({
               message: {
@@ -179,8 +198,10 @@ describe('useTimeline', () => {
           role: MessageRole.Assistant,
           content: '',
           loading: true,
-          canRegenerate: false,
-          canGiveFeedback: false,
+          actions: {
+            canRegenerate: false,
+            canGiveFeedback: false,
+          },
         });
 
         expect(hookResult.result.current.items.length).toBe(3);
@@ -189,8 +210,10 @@ describe('useTimeline', () => {
           role: MessageRole.Assistant,
           content: '',
           loading: true,
-          canRegenerate: false,
-          canGiveFeedback: false,
+          actions: {
+            canRegenerate: false,
+            canGiveFeedback: false,
+          },
         });
 
         act(() => {
@@ -201,8 +224,10 @@ describe('useTimeline', () => {
           role: MessageRole.Assistant,
           content: 'Goodbye',
           loading: true,
-          canRegenerate: false,
-          canGiveFeedback: false,
+          actions: {
+            canRegenerate: false,
+            canGiveFeedback: false,
+          },
         });
 
         act(() => {
@@ -215,8 +240,10 @@ describe('useTimeline', () => {
           role: MessageRole.Assistant,
           content: 'Goodbye',
           loading: false,
-          canRegenerate: true,
-          canGiveFeedback: true,
+          actions: {
+            canRegenerate: true,
+            canGiveFeedback: true,
+          },
         });
       });
 
@@ -240,9 +267,16 @@ describe('useTimeline', () => {
           expect(hookResult.result.current.items.length).toBe(3);
 
           expect(hookResult.result.current.items[2]).toEqual({
-            canEdit: false,
-            canRegenerate: true,
-            canGiveFeedback: false,
+            actions: {
+              canEdit: false,
+              canRegenerate: true,
+              canGiveFeedback: false,
+              canCopy: true,
+            },
+            display: {
+              collapsed: false,
+              hide: false,
+            },
             content: 'My partial',
             id: expect.any(String),
             loading: false,
@@ -262,9 +296,16 @@ describe('useTimeline', () => {
 
           it('updates the last item in the array to be loading', () => {
             expect(hookResult.result.current.items[2]).toEqual({
-              canEdit: false,
-              canRegenerate: false,
-              canGiveFeedback: false,
+              display: {
+                hide: false,
+                collapsed: false,
+              },
+              actions: {
+                canCopy: true,
+                canEdit: false,
+                canRegenerate: false,
+                canGiveFeedback: false,
+              },
               content: '',
               id: expect.any(String),
               loading: true,
@@ -288,9 +329,16 @@ describe('useTimeline', () => {
               expect(hookResult.result.current.items.length).toBe(3);
 
               expect(hookResult.result.current.items[2]).toEqual({
-                canEdit: false,
-                canRegenerate: false,
-                canGiveFeedback: false,
+                actions: {
+                  canCopy: true,
+                  canEdit: false,
+                  canRegenerate: false,
+                  canGiveFeedback: false,
+                },
+                display: {
+                  collapsed: false,
+                  hide: false,
+                },
                 content: '',
                 id: expect.any(String),
                 loading: true,
@@ -308,9 +356,16 @@ describe('useTimeline', () => {
               expect(hookResult.result.current.items.length).toBe(3);
 
               expect(hookResult.result.current.items[2]).toEqual({
-                canEdit: false,
-                canRegenerate: true,
-                canGiveFeedback: true,
+                display: {
+                  collapsed: false,
+                  hide: false,
+                },
+                actions: {
+                  canCopy: true,
+                  canEdit: false,
+                  canRegenerate: true,
+                  canGiveFeedback: true,
+                },
                 content: 'Regenerated',
                 id: expect.any(String),
                 loading: false,
@@ -340,7 +395,7 @@ describe('useTimeline', () => {
             subject.complete();
           });
 
-          props.service.executeFunction.mockResolvedValueOnce({
+          props.chatService.executeFunction.mockResolvedValueOnce({
             content: {
               message: 'my-response',
             },
@@ -364,7 +419,7 @@ describe('useTimeline', () => {
 
           expect(props.onChatComplete).not.toHaveBeenCalled();
 
-          expect(props.service.executeFunction).toHaveBeenCalledWith(
+          expect(props.chatService.executeFunction).toHaveBeenCalledWith(
             'my_function',
             '{}',
             expect.any(Object)

x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts

@@ -7,21 +7,31 @@
 
 import { AbortError } from '@kbn/kibana-utils-plugin/common';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
+import { last } from 'lodash';
 import { useEffect, useMemo, useRef, useState } from 'react';
 import type { Subscription } from 'rxjs';
-import { MessageRole, type ConversationCreateRequest, type Message } from '../../common/types';
+import {
+  ContextDefinition,
+  MessageRole,
+  type ConversationCreateRequest,
+  type Message,
+} from '../../common/types';
 import type { ChatPromptEditorProps } from '../components/chat/chat_prompt_editor';
 import type { ChatTimelineProps } from '../components/chat/chat_timeline';
 import { EMPTY_CONVERSATION_TITLE } from '../i18n';
-import { getSystemMessage } from '../service/get_system_message';
-import type { ObservabilityAIAssistantService, PendingMessage } from '../types';
+import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
+import type { ObservabilityAIAssistantChatService, PendingMessage } from '../types';
 import { getTimelineItemsfromConversation } from '../utils/get_timeline_items_from_conversation';
 import type { UseGenAIConnectorsResult } from './use_genai_connectors';
 
-export function createNewConversation(): ConversationCreateRequest {
+export function createNewConversation({
+  contexts,
+}: {
+  contexts: ContextDefinition[];
+}): ConversationCreateRequest {
   return {
     '@timestamp': new Date().toISOString(),
-    messages: [getSystemMessage()],
+    messages: [getAssistantSetupMessage({ contexts })],
     conversation: {
       title: EMPTY_CONVERSATION_TITLE,
     },
@@ -41,14 +51,14 @@ export function useTimeline({
   messages,
   connectors,
   currentUser,
-  service,
+  chatService,
   onChatUpdate,
   onChatComplete,
 }: {
   messages: Message[];
   connectors: UseGenAIConnectorsResult;
   currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
-  service: ObservabilityAIAssistantService;
+  chatService: ObservabilityAIAssistantChatService;
   onChatUpdate: (messages: Message[]) => void;
   onChatComplete: (messages: Message[]) => void;
 }): UseTimelineResult {
@@ -57,12 +67,15 @@ export function useTimeline({
   const hasConnector = !!connectorId;
 
   const conversationItems = useMemo(() => {
-    return getTimelineItemsfromConversation({
+    const items = getTimelineItemsfromConversation({
       messages,
       currentUser,
       hasConnector,
+      chatService,
     });
-  }, [messages, currentUser, hasConnector]);
+
+    return items;
+  }, [messages, currentUser, hasConnector, chatService]);
 
   const [subscription, setSubscription] = useState<Subscription | undefined>();
 
@@ -73,7 +86,7 @@ export function useTimeline({
   function chat(nextMessages: Message[]): Promise<Message[]> {
     const controller = new AbortController();
 
-    return new Promise<PendingMessage>((resolve, reject) => {
+    return new Promise<PendingMessage | undefined>((resolve, reject) => {
       if (!connectorId) {
         reject(new Error('Can not add a message without a connector'));
         return;
@@ -81,7 +94,18 @@ export function useTimeline({
 
       onChatUpdate(nextMessages);
 
-      const response$ = service.chat({ messages: nextMessages, connectorId });
+      const lastMessage = last(nextMessages);
+
+      if (lastMessage?.message.function_call?.name) {
+        // the user has edited a function suggestion, no need to talk to
+        resolve(undefined);
+        return;
+      }
+
+      const response$ = chatService!.chat({
+        messages: nextMessages,
+        connectorId,
+      });
 
       let pendingMessageLocal = pendingMessage;
 
@@ -101,31 +125,35 @@ export function useTimeline({
         return nextSubscription;
       });
     }).then(async (reply) => {
-      if (reply.error) {
+      if (reply?.error) {
         return nextMessages;
       }
-      if (reply.aborted) {
+      if (reply?.aborted) {
         return nextMessages;
       }
 
       setPendingMessage(undefined);
 
-      const messagesAfterChat = nextMessages.concat({
-        '@timestamp': new Date().toISOString(),
-        message: {
-          ...reply.message,
-        },
-      });
+      const messagesAfterChat = reply
+        ? nextMessages.concat({
+            '@timestamp': new Date().toISOString(),
+            message: {
+              ...reply.message,
+            },
+          })
+        : nextMessages;
 
       onChatUpdate(messagesAfterChat);
 
-      if (reply?.message.function_call?.name) {
-        const name = reply.message.function_call.name;
+      const lastMessage = last(messagesAfterChat);
+
+      if (lastMessage?.message.function_call?.name) {
+        const name = lastMessage.message.function_call.name;
 
         try {
-          const message = await service.executeFunction(
+          const message = await chatService!.executeFunction(
             name,
-            reply.message.function_call.arguments,
+            lastMessage.message.function_call.arguments,
             controller.signal
           );
 
@@ -133,8 +161,8 @@ export function useTimeline({
             messagesAfterChat.concat({
               '@timestamp': new Date().toISOString(),
               message: {
-                role: MessageRole.User,
                 name,
+                role: MessageRole.User,
                 content: JSON.stringify(message.content),
                 data: JSON.stringify(message.data),
               },
@@ -149,7 +177,7 @@ export function useTimeline({
                 name,
                 content: JSON.stringify({
                   message: error.toString(),
-                  ...error.body,
+                  error: error.body,
                 }),
               },
             })
@@ -165,16 +193,23 @@ export function useTimeline({
     if (pendingMessage) {
       return conversationItems.concat({
         id: '',
-        canEdit: false,
-        canRegenerate: pendingMessage.aborted || !!pendingMessage.error,
-        canGiveFeedback: false,
-        title: '',
-        role: pendingMessage.message.role,
+        actions: {
+          canCopy: true,
+          canEdit: false,
+          canGiveFeedback: false,
+          canRegenerate: pendingMessage.aborted || !!pendingMessage.error,
+        },
+        display: {
+          collapsed: false,
+          hide: pendingMessage.message.role === MessageRole.System,
+        },
         content: pendingMessage.message.content,
-        loading: !pendingMessage.aborted && !pendingMessage.error,
-        function_call: pendingMessage.message.function_call,
         currentUser,
         error: pendingMessage.error,
+        function_call: pendingMessage.message.function_call,
+        loading: !pendingMessage.aborted && !pendingMessage.error,
+        role: pendingMessage.message.role,
+        title: '',
       });
     }
 
@@ -189,7 +224,12 @@ export function useTimeline({
 
   return {
     items,
-    onEdit: (item, content) => {},
+    onEdit: async (item, newMessage) => {
+      const indexOf = items.indexOf(item);
+      const sliced = messages.slice(0, indexOf - 1);
+      const nextMessages = await chat(sliced.concat(newMessage));
+      onChatComplete(nextMessages);
+    },
     onFeedback: (item, feedback) => {},
     onRegenerate: (item) => {
       const indexOf = items.indexOf(item);

x-pack/plugins/observability_ai_assistant/public/plugin.tsx

@@ -17,12 +17,6 @@ import { i18n } from '@kbn/i18n';
 import type { Logger } from '@kbn/logging';
 import React from 'react';
 import ReactDOM from 'react-dom';
-import type {
-  ContextRegistry,
-  FunctionRegistry,
-  RegisterContextDefinition,
-  RegisterFunctionDefinition,
-} from '../common/types';
 import { registerFunctions } from './functions';
 import { createService } from './service/create_service';
 import type {
@@ -101,35 +95,22 @@ export class ObservabilityAIAssistantPlugin
     coreStart: CoreStart,
     pluginsStart: ObservabilityAIAssistantPluginStartDependencies
   ): ObservabilityAIAssistantPluginStart {
-    const contextRegistry: ContextRegistry = new Map();
-    const functionRegistry: FunctionRegistry = new Map();
-
     const service = (this.service = createService({
       coreStart,
       securityStart: pluginsStart.security,
-      contextRegistry,
-      functionRegistry,
       enabled: coreStart.application.capabilities.observabilityAIAssistant.show === true,
     }));
 
-    const registerContext: RegisterContextDefinition = (context) => {
-      contextRegistry.set(context.name, context);
-    };
-
-    const registerFunction: RegisterFunctionDefinition = (def, respond, render) => {
-      functionRegistry.set(def.name, { options: def, respond, render });
-    };
-
-    registerFunctions({
-      registerContext,
-      registerFunction,
-      service,
+    service.register(({ signal, registerContext, registerFunction }) => {
+      return registerFunctions({
+        service,
+        signal,
+        pluginsStart,
+        registerContext,
+        registerFunction,
+      });
     });
 
-    return {
-      ...service,
-      registerContext,
-      registerFunction,
-    };
+    return service;
   }
 }

x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx

@@ -10,12 +10,14 @@ import { i18n } from '@kbn/i18n';
 import React, { useMemo, useState } from 'react';
 import { ChatBody } from '../../components/chat/chat_body';
 import { ConversationList } from '../../components/chat/conversation_list';
+import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
 import { useAbortableAsync } from '../../hooks/use_abortable_async';
 import { useConfirmModal } from '../../hooks/use_confirm_modal';
 import { useConversation } from '../../hooks/use_conversation';
 import { useCurrentUser } from '../../hooks/use_current_user';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
 import { useKibana } from '../../hooks/use_kibana';
+import { useKnowledgeBase } from '../../hooks/use_knowledge_base';
 import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import { useObservabilityAIAssistantParams } from '../../hooks/use_observability_ai_assistant_params';
 import { useObservabilityAIAssistantRouter } from '../../hooks/use_observability_ai_assistant_router';
@@ -33,6 +35,8 @@ const chatBodyContainerClassNameWithError = css`
 export function ConversationView() {
   const connectors = useGenAIConnectors();
 
+  const knowledgeBase = useKnowledgeBase();
+
   const currentUser = useCurrentUser();
 
   const service = useObservabilityAIAssistant();
@@ -59,10 +63,19 @@ export function ConversationView() {
 
   const [isUpdatingList, setIsUpdatingList] = useState(false);
 
+  const chatService = useAbortableAsync(
+    ({ signal }) => {
+      return service.start({ signal });
+    },
+    [service]
+  );
+
   const conversationId = 'conversationId' in path ? path.conversationId : undefined;
 
-  const { conversation, displayedMessages, setDisplayedMessages, save } =
-    useConversation(conversationId);
+  const { conversation, displayedMessages, setDisplayedMessages, save } = useConversation({
+    conversationId,
+    chatService: chatService.value,
+  });
 
   const conversations = useAbortableAsync(
     ({ signal }) => {
@@ -173,6 +186,7 @@ export function ConversationView() {
                 });
             }}
           />
+          <EuiSpacer size="m" />
         </EuiFlexItem>
         <EuiFlexItem
           grow
@@ -196,29 +210,38 @@ export function ConversationView() {
               })}
             </EuiCallOut>
           ) : null}
-          {conversation.loading ? <EuiLoadingSpinner /> : null}
-          {!conversation.error && conversation.value ? (
-            <ChatBody
-              currentUser={currentUser}
-              connectors={connectors}
-              title={conversation.value.conversation.title}
-              connectorsManagementHref={getConnectorsManagementHref(http)}
-              service={service}
-              messages={displayedMessages}
-              onChatComplete={(messages) => {
-                save(messages)
-                  .then((nextConversation) => {
-                    conversations.refresh();
-                    if (!conversationId) {
-                      navigateToConversation(nextConversation.conversation.id);
-                    }
-                  })
-                  .catch(() => {});
-              }}
-              onChatUpdate={(messages) => {
-                setDisplayedMessages(messages);
-              }}
-            />
+          {chatService.loading || conversation.loading ? (
+            <EuiFlexGroup direction="column" alignItems="center" gutterSize="l">
+              <EuiFlexItem grow={false}>
+                <EuiSpacer size="xl" />
+                <EuiLoadingSpinner size="l" />
+              </EuiFlexItem>
+            </EuiFlexGroup>
+          ) : null}
+          {!conversation.error && conversation.value && chatService.value ? (
+            <ObservabilityAIAssistantChatServiceProvider value={chatService.value}>
+              <ChatBody
+                currentUser={currentUser}
+                connectors={connectors}
+                knowledgeBase={knowledgeBase}
+                title={conversation.value.conversation.title}
+                connectorsManagementHref={getConnectorsManagementHref(http)}
+                messages={displayedMessages}
+                onChatComplete={(messages) => {
+                  save(messages)
+                    .then((nextConversation) => {
+                      conversations.refresh();
+                      if (!conversationId) {
+                        navigateToConversation(nextConversation.conversation.id);
+                      }
+                    })
+                    .catch(() => {});
+                }}
+                onChatUpdate={(messages) => {
+                  setDisplayedMessages(messages);
+                }}
+              />
+            </ObservabilityAIAssistantChatServiceProvider>
           ) : null}
           <EuiSpacer size="m" />
         </EuiFlexItem>

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

@@ -4,19 +4,17 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import type { CoreStart, HttpFetchOptions } from '@kbn/core/public';
-import { ReadableStream } from 'stream/web';
-import type { AuthenticatedUser } from '@kbn/security-plugin/common';
-import type { ObservabilityAIAssistantService } from '../types';
-import { createService } from './create_service';
-import { SecurityPluginStart } from '@kbn/security-plugin/public';
+import type { HttpFetchOptions } from '@kbn/core/public';
 import { lastValueFrom } from 'rxjs';
+import { ReadableStream } from 'stream/web';
+import type { ObservabilityAIAssistantChatService } from '../types';
+import { createChatService } from './create_chat_service';
 
-describe('createService', () => {
+describe('createChatService', () => {
   describe('chat', () => {
-    let service: ObservabilityAIAssistantService;
+    let service: ObservabilityAIAssistantChatService;
 
-    const httpPostSpy = jest.fn();
+    const clientSpy = jest.fn();
 
     function respondWithChunks({ chunks, status = 200 }: { status?: number; chunks: string[] }) {
       const response = {
@@ -33,33 +31,23 @@ describe('createService', () => {
         },
       };
 
-      httpPostSpy.mockResolvedValueOnce(response);
+      clientSpy.mockResolvedValueOnce(response);
     }
 
     function chat() {
       return service.chat({ messages: [], connectorId: '' });
     }
 
-    beforeEach(() => {
-      service = createService({
-        coreStart: {
-          http: {
-            post: httpPostSpy,
-          },
-        } as unknown as CoreStart,
-        securityStart: {
-          authc: {
-            getCurrentUser: () => Promise.resolve({ username: 'elastic' } as AuthenticatedUser),
-          },
-        } as unknown as SecurityPluginStart,
-        contextRegistry: new Map(),
-        functionRegistry: new Map(),
-        enabled: true,
+    beforeEach(async () => {
+      service = await createChatService({
+        client: clientSpy,
+        registrations: [],
+        signal: new AbortController().signal,
       });
     });
 
     afterEach(() => {
-      httpPostSpy.mockReset();
+      clientSpy.mockReset();
     });
 
     it('correctly parses a stream of JSON lines', async () => {
@@ -169,7 +157,7 @@ describe('createService', () => {
     });
 
     it('cancels a running http request when aborted', async () => {
-      httpPostSpy.mockImplementationOnce((endpoint: string, options: HttpFetchOptions) => {
+      clientSpy.mockImplementationOnce((endpoint: string, options: HttpFetchOptions) => {
         options.signal?.addEventListener('abort', () => {
           expect(options.signal?.aborted).toBeTruthy();
         });

x-pack/plugins/observability_ai_assistant/public/service/create_chat_service.ts

@@ -0,0 +1,226 @@
+/*
+ * 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 { IncomingMessage } from 'http';
+import { cloneDeep, pick } from 'lodash';
+import {
+  BehaviorSubject,
+  map,
+  filter as rxJsFilter,
+  scan,
+  catchError,
+  of,
+  concatMap,
+  shareReplay,
+  finalize,
+  delay,
+} from 'rxjs';
+import { HttpResponse } from '@kbn/core/public';
+import { AbortError } from '@kbn/kibana-utils-plugin/common';
+import {
+  type RegisterContextDefinition,
+  type RegisterFunctionDefinition,
+  Message,
+  MessageRole,
+  ContextRegistry,
+  FunctionRegistry,
+} from '../../common/types';
+import { ObservabilityAIAssistantAPIClient } from '../api';
+import type {
+  ChatRegistrationFunction,
+  CreateChatCompletionResponseChunk,
+  ObservabilityAIAssistantChatService,
+  PendingMessage,
+} from '../types';
+import { readableStreamReaderIntoObservable } from '../utils/readable_stream_reader_into_observable';
+
+export async function createChatService({
+  signal: setupAbortSignal,
+  registrations,
+  client,
+}: {
+  signal: AbortSignal;
+  registrations: ChatRegistrationFunction[];
+  client: ObservabilityAIAssistantAPIClient;
+}): Promise<ObservabilityAIAssistantChatService> {
+  const contextRegistry: ContextRegistry = new Map();
+  const functionRegistry: FunctionRegistry = new Map();
+
+  const registerContext: RegisterContextDefinition = (context) => {
+    contextRegistry.set(context.name, context);
+  };
+
+  const registerFunction: RegisterFunctionDefinition = (def, respond, render) => {
+    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 }))
+  );
+
+  return {
+    executeFunction: async (name, args, signal) => {
+      const fn = functionRegistry.get(name);
+
+      if (!fn) {
+        throw new Error(`Function ${name} not found`);
+      }
+
+      const parsedArguments = args ? JSON.parse(args) : {};
+
+      // validate
+
+      return await fn.respond({ arguments: parsedArguments }, signal);
+    },
+    renderFunction: (name, args, response) => {
+      const fn = functionRegistry.get(name);
+
+      if (!fn) {
+        throw new Error(`Function ${name} not found`);
+      }
+
+      const parsedArguments = args ? JSON.parse(args) : {};
+
+      const parsedResponse = {
+        content: JSON.parse(response.content ?? '{}'),
+        data: JSON.parse(response.data ?? '{}'),
+      };
+
+      // validate
+
+      return fn.render?.({ response: parsedResponse, arguments: parsedArguments });
+    },
+    getContexts,
+    getFunctions,
+    hasRenderFunction: (name: string) => {
+      return !!getFunctions().find((fn) => fn.options.name === name)?.render;
+    },
+    chat({ connectorId, messages }: { connectorId: string; messages: Message[] }) {
+      const subject = new BehaviorSubject<PendingMessage>({
+        message: {
+          role: MessageRole.Assistant,
+        },
+      });
+
+      const contexts = ['core', 'apm'];
+
+      const functions = getFunctions({ contexts });
+
+      const controller = new AbortController();
+
+      client('POST /internal/observability_ai_assistant/chat', {
+        params: {
+          body: {
+            messages,
+            connectorId,
+            functions: functions.map((fn) => pick(fn.options, 'name', 'description', 'parameters')),
+          },
+        },
+        signal: controller.signal,
+        asResponse: true,
+        rawResponse: true,
+      })
+        .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)
+            .pipe(
+              map((line) => line.substring(6)),
+              rxJsFilter((line) => !!line && line !== '[DONE]'),
+              map((line) => JSON.parse(line) as CreateChatCompletionResponseChunk),
+              rxJsFilter((line) => line.object === 'chat.completion.chunk'),
+              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,
+                  },
+                }
+              ),
+              catchError((error) =>
+                of({
+                  ...subject.value,
+                  error,
+                  aborted: error instanceof AbortError || controller.signal.aborted,
+                })
+              )
+            )
+            .subscribe(subject);
+
+          controller.signal.addEventListener('abort', () => {
+            subscription.unsubscribe();
+            subject.next({
+              ...subject.value,
+              aborted: true,
+            });
+            subject.complete();
+          });
+        })
+        .catch((err) => {
+          subject.next({
+            ...subject.value,
+            aborted: false,
+            error: err,
+          });
+          subject.complete();
+        });
+
+      return subject.pipe(
+        concatMap((value) => of(value).pipe(delay(50))),
+        shareReplay(1),
+        finalize(() => {
+          controller.abort();
+        })
+      );
+    },
+  };
+}

x-pack/plugins/observability_ai_assistant/public/service/create_service.ts

@@ -5,200 +5,37 @@
  * 2.0.
  */
 
-import type { CoreStart, HttpResponse } from '@kbn/core/public';
-import { AbortError } from '@kbn/kibana-utils-plugin/common';
+import type { CoreStart } from '@kbn/core/public';
 import { SecurityPluginStart } from '@kbn/security-plugin/public';
-import { IncomingMessage } from 'http';
-import { cloneDeep } from 'lodash';
-import {
-  BehaviorSubject,
-  catchError,
-  concatMap,
-  delay,
-  filter as rxJsFilter,
-  finalize,
-  map,
-  of,
-  scan,
-  shareReplay,
-} from 'rxjs';
-import type { Message } from '../../common';
-import { ContextRegistry, FunctionRegistry, MessageRole } from '../../common/types';
 import { createCallObservabilityAIAssistantAPI } from '../api';
-import type {
-  CreateChatCompletionResponseChunk,
-  ObservabilityAIAssistantService,
-  PendingMessage,
-} from '../types';
-import { readableStreamReaderIntoObservable } from '../utils/readable_stream_reader_into_observable';
+import type { ChatRegistrationFunction, ObservabilityAIAssistantService } from '../types';
+import { createChatService } from './create_chat_service';
 
 export function createService({
   coreStart,
   securityStart,
-  functionRegistry,
-  contextRegistry,
   enabled,
 }: {
   coreStart: CoreStart;
   securityStart: SecurityPluginStart;
-  functionRegistry: FunctionRegistry;
-  contextRegistry: ContextRegistry;
   enabled: boolean;
-}): ObservabilityAIAssistantService {
+}): ObservabilityAIAssistantService & { register: (fn: ChatRegistrationFunction) => void } {
   const client = createCallObservabilityAIAssistantAPI(coreStart);
 
-  const getContexts: ObservabilityAIAssistantService['getContexts'] = () => {
-    return Array.from(contextRegistry.values());
-  };
-  const getFunctions: ObservabilityAIAssistantService['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;
-  };
+  const registrations: ChatRegistrationFunction[] = [];
 
   return {
     isEnabled: () => {
       return enabled;
     },
-    chat({ connectorId, messages }: { connectorId: string; messages: Message[] }) {
-      const subject = new BehaviorSubject<PendingMessage>({
-        message: {
-          role: MessageRole.Assistant,
-        },
-      });
-
-      const contexts = ['core'];
-
-      const functions = getFunctions({ contexts });
-
-      const controller = new AbortController();
-
-      client('POST /internal/observability_ai_assistant/chat', {
-        params: {
-          body: {
-            messages,
-            connectorId,
-            functions: functions.map((fn) => fn.options),
-          },
-        },
-        signal: controller.signal,
-        asResponse: true,
-        rawResponse: true,
-      })
-        .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)
-            .pipe(
-              map((line) => line.substring(6)),
-              rxJsFilter((line) => !!line && line !== '[DONE]'),
-              map((line) => JSON.parse(line) as CreateChatCompletionResponseChunk),
-              rxJsFilter((line) => line.object === 'chat.completion.chunk'),
-              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,
-                  },
-                }
-              ),
-              catchError((error) =>
-                of({
-                  ...subject.value,
-                  error,
-                  aborted: error instanceof AbortError || controller.signal.aborted,
-                })
-              )
-            )
-            .subscribe(subject);
-
-          controller.signal.addEventListener('abort', () => {
-            subscription.unsubscribe();
-            subject.next({
-              ...subject.value,
-              aborted: true,
-            });
-            subject.complete();
-          });
-        })
-        .catch((err) => {
-          subject.next({
-            ...subject.value,
-            aborted: false,
-            error: err,
-          });
-          subject.complete();
-        });
-
-      return subject.pipe(
-        concatMap((value) => of(value).pipe(delay(50))),
-        shareReplay(1),
-        finalize(() => {
-          controller.abort();
-        })
-      );
+    register: (fn) => {
+      registrations.push(fn);
     },
+    start: async ({ signal }) => {
+      return await createChatService({ client, signal, registrations });
+    },
+
     callApi: client,
     getCurrentUser: () => securityStart.authc.getCurrentUser(),
-    getContexts,
-    getFunctions,
-    executeFunction: async (name, args, signal) => {
-      const fn = functionRegistry.get(name);
-
-      if (!fn) {
-        throw new Error(`Function ${name} not found`);
-      }
-
-      const parsedArguments = args ? JSON.parse(args) : {};
-
-      // validate
-
-      return await fn.respond({ arguments: parsedArguments }, signal);
-    },
-    renderFunction: (name, response) => {
-      const fn = functionRegistry.get(name);
-
-      if (!fn) {
-        throw new Error(`Function ${name} not found`);
-      }
-
-      return fn.render?.({ response });
-    },
   };
 }

x-pack/plugins/observability_ai_assistant/public/service/get_assistant_setup_message.ts

@@ -0,0 +1,26 @@
+/*
+ * 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 { MessageRole } from '../../common';
+import { ContextDefinition } from '../../common/types';
+
+export function getAssistantSetupMessage({ contexts }: { contexts: ContextDefinition[] }) {
+  return {
+    '@timestamp': new Date().toISOString(),
+    message: {
+      role: MessageRole.System as const,
+      content: [
+        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`
+        ),
+      ]
+        .concat(contexts.map((context) => context.description))
+        .join('\n'),
+    },
+  };
+}

x-pack/plugins/observability_ai_assistant/public/service/get_system_message.ts

@@ -1,47 +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 dedent from 'dedent';
-import { MessageRole } from '../../common';
-
-export function getSystemMessage() {
-  return {
-    '@timestamp': new Date().toISOString(),
-    message: {
-      role: MessageRole.System as const,
-      content:
-        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.
-
-  You can use the "summarise" 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. When you create this summarisation, make sure you craft it in a way that can be recalled with a semantic search later.
-
-  Additionally, you can use the "recall" function to retrieve relevant information from the knowledge database. Using the "recall" function will allow you to create a much more tailored and improved user experience. Assume the user has told you before.
-
-  For every question that the user asks, you, as the assistant, must query the knowledge base, no matter how confident you are in your response.
-
-  When the user starts a conversation, you as the assistant, MUST ALWAYS use the "recall" function to make sure you have enough context from the knowledge database to understand their question. DO NOT assume your knowledge is sufficient to answer the question appropriately. Make sure you mention this to the user in your first reply to them.
-
-  You do have access to private data that the user has shared with you in previous conversations.
-
-  Given this ability to store and recall data, it is very important to not make assumptions about the user's data or preferences, but rather, first query the knowledge database, and if nothing relevant comes up, ask the user for clarification. This is very important! They could be storing their data in any field, and in any data stream or index.
-
-  RIGHT:
-  User: "What is NASA"
-  Assistant executes recall function
-  Assistant answers question with data from recall function response
-
-  WRONG:
-  User: "What is NASA"
-  Assistant answers question without querying the knowledge
-
-  You should autonomously execute these functions - do not wait on the user's permission, but be proactive.
-
-  Note that any visualisations will be displayed ABOVE your textual response, not below.
-
-  Feel free to use Markdown in your replies, especially for code and query statements.`),
-    },
-  };
-}

x-pack/plugins/observability_ai_assistant/public/types.ts

@@ -49,13 +49,11 @@ export interface PendingMessage {
   error?: any;
 }
 
-export interface ObservabilityAIAssistantService {
-  isEnabled: () => boolean;
+export interface ObservabilityAIAssistantChatService {
   chat: (options: { messages: Message[]; connectorId: string }) => Observable<PendingMessage>;
-  callApi: ObservabilityAIAssistantAPIClient;
-  getCurrentUser: () => Promise<AuthenticatedUser>;
   getContexts: () => ContextDefinition[];
   getFunctions: (options?: { contexts?: string[]; filter?: string }) => FunctionDefinition[];
+  hasRenderFunction: (name: string) => boolean;
   executeFunction: (
     name: string,
     args: string | undefined,
@@ -63,13 +61,26 @@ export interface ObservabilityAIAssistantService {
   ) => Promise<{ content?: Serializable; data?: Serializable }>;
   renderFunction: (
     name: string,
-    response: { data?: Serializable; content?: Serializable }
+    args: string | undefined,
+    response: { data?: string; content?: string }
   ) => React.ReactNode;
 }
 
-export interface ObservabilityAIAssistantPluginStart extends ObservabilityAIAssistantService {
-  registerContext: RegisterContextDefinition;
+export type ChatRegistrationFunction = ({}: {
+  signal: AbortSignal;
   registerFunction: RegisterFunctionDefinition;
+  registerContext: RegisterContextDefinition;
+}) => Promise<void>;
+
+export interface ObservabilityAIAssistantService {
+  isEnabled: () => boolean;
+  callApi: ObservabilityAIAssistantAPIClient;
+  getCurrentUser: () => Promise<AuthenticatedUser>;
+  start: ({}: { signal: AbortSignal }) => Promise<ObservabilityAIAssistantChatService>;
+}
+
+export interface ObservabilityAIAssistantPluginStart extends ObservabilityAIAssistantService {
+  register: (fn: ChatRegistrationFunction) => void;
 }
 
 export interface ObservabilityAIAssistantPluginSetup {}

x-pack/plugins/observability_ai_assistant/public/utils/builders.ts

@@ -5,26 +5,39 @@
  * 2.0.
  */
 
-import { uniqueId } from 'lodash';
+import { merge, uniqueId } from 'lodash';
 import { MessageRole, Conversation, FunctionDefinition } from '../../common/types';
 import { ChatTimelineItem } from '../components/chat/chat_timeline';
-import { getSystemMessage } from '../service/get_system_message';
+import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
 
-type ChatItemBuildProps = Partial<ChatTimelineItem> & Pick<ChatTimelineItem, 'role'>;
+type ChatItemBuildProps = Omit<Partial<ChatTimelineItem>, 'actions' | 'display' | 'currentUser'> & {
+  actions?: Partial<ChatTimelineItem['actions']>;
+  display?: Partial<ChatTimelineItem['display']>;
+  currentUser?: Partial<ChatTimelineItem['currentUser']>;
+} & Pick<ChatTimelineItem, 'role'>;
 
 export function buildChatItem(params: ChatItemBuildProps): ChatTimelineItem {
-  return {
-    id: uniqueId(),
-    title: '',
-    canEdit: false,
-    canGiveFeedback: false,
-    canRegenerate: params.role === MessageRole.User,
-    currentUser: {
-      username: 'elastic',
+  return merge(
+    {
+      id: uniqueId(),
+      title: '',
+      actions: {
+        canCopy: true,
+        canEdit: false,
+        canGiveFeedback: false,
+        canRegenerate: params.role === MessageRole.Assistant,
+      },
+      display: {
+        collapsed: false,
+        hide: false,
+      },
+      currentUser: {
+        username: 'elastic',
+      },
+      loading: false,
     },
-    loading: false,
-    ...params,
-  };
+    params
+  );
 }
 
 export function buildSystemChatItem(params?: Omit<ChatItemBuildProps, 'role'>) {
@@ -38,7 +51,12 @@ export function buildChatInitItem() {
   return buildChatItem({
     role: MessageRole.User,
     title: 'started a conversation',
-    canRegenerate: false,
+    actions: {
+      canEdit: false,
+      canCopy: true,
+      canGiveFeedback: false,
+      canRegenerate: false,
+    },
   });
 }
 
@@ -46,7 +64,12 @@ export function buildUserChatItem(params?: Omit<ChatItemBuildProps, 'role'>) {
   return buildChatItem({
     role: MessageRole.User,
     content: "What's a function?",
-    canEdit: true,
+    actions: {
+      canCopy: true,
+      canEdit: true,
+      canGiveFeedback: false,
+      canRegenerate: true,
+    },
     ...params,
   });
 }
@@ -56,8 +79,12 @@ export function buildAssistantChatItem(params?: Omit<ChatItemBuildProps, 'role'>
     role: MessageRole.Assistant,
     content: `In computer programming and mathematics, a function is a fundamental concept that represents a relationship between input values and output values. It takes one or more input values (also known as arguments or parameters) and processes them to produce a result, which is the output of the function. The input values are passed to the function, and the function performs a specific set of operations or calculations on those inputs to produce the desired output.
     A function is often defined with a name, which serves as an identifier to call and use the function in the code. It can be thought of as a reusable block of code that can be executed whenever needed, and it helps in organizing code and making it more modular and maintainable.`,
-    canRegenerate: true,
-    canGiveFeedback: true,
+    actions: {
+      canCopy: true,
+      canEdit: false,
+      canRegenerate: true,
+      canGiveFeedback: true,
+    },
     ...params,
   });
 }
@@ -92,7 +119,7 @@ export function buildConversation(params?: Partial<Conversation>) {
       title: '',
       last_updated: '',
     },
-    messages: [getSystemMessage()],
+    messages: [getAssistantSetupMessage({ contexts: [] })],
     labels: {},
     numeric_labels: {},
     namespace: '',
@@ -106,6 +133,7 @@ export function buildFunction(): FunctionDefinition {
       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: {
@@ -135,6 +163,7 @@ export function buildFunctionServiceSummary(): FunctionDefinition {
       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',
       },

x-pack/plugins/observability_ai_assistant/public/utils/create_initialized_object.test.ts

@@ -0,0 +1,134 @@
+/*
+ * 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 { createInitializedObject } from './create_initialized_object';
+
+describe('createInitializedObject', () => {
+  it('should return an object with properties of type "string" set to a default value of ""', () => {
+    expect(
+      createInitializedObject({
+        type: 'object',
+        properties: {
+          foo: {
+            type: 'string',
+          },
+        },
+        required: ['foo'],
+      })
+    ).toStrictEqual({ foo: '' });
+  });
+
+  it('should return an object with properties of type "number" set to a default value of 1', () => {
+    expect(
+      createInitializedObject({
+        type: 'object',
+        properties: {
+          foo: {
+            type: 'number',
+          },
+        },
+        required: ['foo'],
+      })
+    ).toStrictEqual({ foo: 1 });
+  });
+
+  it('should return an object with properties of type "array" set to a default value of []', () => {
+    expect(
+      createInitializedObject({
+        type: 'object',
+        properties: {
+          foo: {
+            type: 'array',
+          },
+        },
+        required: ['foo'],
+      })
+    ).toStrictEqual({ foo: [] });
+  });
+
+  it('should return an object with default values for properties that are required', () => {
+    expect(
+      createInitializedObject({
+        type: 'object',
+        properties: {
+          requiredProperty: {
+            type: 'string',
+          },
+          notRequiredProperty: {
+            type: 'string',
+          },
+        },
+        required: ['requiredProperty'],
+      })
+    ).toStrictEqual({ requiredProperty: '' });
+  });
+
+  it('should return an object with nested fields if they are present in the schema', () => {
+    expect(
+      createInitializedObject({
+        type: 'object',
+        properties: {
+          foo: {
+            type: 'object',
+            properties: {
+              bar: {
+                type: 'object',
+                properties: {
+                  baz: {
+                    type: 'string',
+                  },
+                },
+                required: ['baz'],
+              },
+            },
+            required: ['bar'],
+          },
+        },
+      })
+    ).toStrictEqual({ foo: { bar: { baz: '' } } });
+
+    expect(
+      createInitializedObject({
+        type: 'object',
+        properties: {
+          foo: {
+            type: 'object',
+            properties: {
+              bar: {
+                type: 'string',
+              },
+            },
+          },
+        },
+      })
+    ).toStrictEqual({ foo: {} });
+  });
+
+  it('should handle a real life example', () => {
+    expect(
+      createInitializedObject({
+        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',
+          },
+          notRequired: {
+            type: 'string',
+            description: 'This property is not required.',
+          },
+        },
+        required: ['method', 'path'] as const,
+      })
+    ).toStrictEqual({ method: '', path: '' });
+  });
+});

x-pack/plugins/observability_ai_assistant/public/utils/create_initialized_object.ts

@@ -0,0 +1,42 @@
+/*
+ * 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 '../../common/types';
+
+type Params = FunctionDefinition['options']['parameters'];
+
+export function createInitializedObject(parameters: Params) {
+  const emptyObject: Record<string, string | any> = {};
+
+  function traverseProperties({ properties, required }: Params) {
+    for (const propName in properties) {
+      if (properties.hasOwnProperty(propName)) {
+        const prop = properties[propName] as Params;
+
+        if (prop.type === 'object') {
+          emptyObject[propName] = createInitializedObject(prop);
+        } else if (required?.includes(propName)) {
+          if (prop.type === 'array') {
+            emptyObject[propName] = [];
+          }
+
+          if (prop.type === 'number') {
+            emptyObject[propName] = 1;
+          }
+
+          if (prop.type === 'string') {
+            emptyObject[propName] = '';
+          }
+        }
+      }
+    }
+  }
+
+  traverseProperties(parameters);
+
+  return emptyObject;
+}

x-pack/plugins/observability_ai_assistant/public/utils/get_role_translation.ts

@@ -0,0 +1,30 @@
+/*
+ * 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 { MessageRole } from '../../common';
+
+export function getRoleTranslation(role: MessageRole) {
+  if (role === MessageRole.User) {
+    return i18n.translate('xpack.observabilityAiAssistant.chatTimeline.messages.user.label', {
+      defaultMessage: 'You',
+    });
+  }
+
+  if (role === MessageRole.System) {
+    return i18n.translate('xpack.observabilityAiAssistant.chatTimeline.messages.system.label', {
+      defaultMessage: 'System',
+    });
+  }
+
+  return i18n.translate(
+    'xpack.observabilityAiAssistant.chatTimeline.messages.elasticAssistant.label',
+    {
+      defaultMessage: 'Elastic Assistant',
+    }
+  );
+}

x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.ts

@@ -1,79 +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 { v4 } from 'uuid';
-import { i18n } from '@kbn/i18n';
-import type { AuthenticatedUser } from '@kbn/security-plugin/common';
-import dedent from 'dedent';
-import { type Message, MessageRole } from '../../common';
-import type { ChatTimelineItem } from '../components/chat/chat_timeline';
-
-export function getTimelineItemsfromConversation({
-  currentUser,
-  messages,
-  hasConnector,
-}: {
-  currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
-  messages: Message[];
-  hasConnector: boolean;
-}): ChatTimelineItem[] {
-  return [
-    {
-      id: v4(),
-      role: MessageRole.User,
-      title: i18n.translate('xpack.observabilityAiAssistant.conversationStartTitle', {
-        defaultMessage: 'started a conversation',
-      }),
-      canEdit: false,
-      canGiveFeedback: false,
-      canRegenerate: false,
-      loading: false,
-      currentUser,
-    },
-    ...messages.map((message) => {
-      const hasFunction = !!message.message.function_call?.name;
-      const isSystemPrompt = message.message.role === MessageRole.System;
-
-      let title: string;
-      let content: string | undefined;
-
-      if (hasFunction) {
-        title = i18n.translate('xpack.observabilityAiAssistant.suggestedFunctionEvent', {
-          defaultMessage: 'suggested a function',
-        });
-        content = dedent(`I have requested your system performs the function _${
-          message.message.function_call?.name
-        }_ with the payload 
-        \`\`\`
-        ${JSON.stringify(JSON.parse(message.message.function_call?.arguments || ''), null, 4)}
-        \`\`\`
-        and return its results for me to look at.`);
-      } else if (isSystemPrompt) {
-        title = i18n.translate('xpack.observabilityAiAssistant.addedSystemPromptEvent', {
-          defaultMessage: 'added a prompt',
-        });
-        content = '';
-      } else {
-        title = '';
-        content = message.message.content;
-      }
-
-      const props = {
-        id: v4(),
-        role: message.message.role,
-        canEdit: hasConnector && (message.message.role === MessageRole.User || hasFunction),
-        canRegenerate: hasConnector && message.message.role === MessageRole.Assistant,
-        canGiveFeedback: message.message.role === MessageRole.Assistant,
-        loading: false,
-        title,
-        content,
-        currentUser,
-      };
-
-      return props;
-    }),
-  ];
-}

x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.tsx

@@ -0,0 +1,222 @@
+/*
+ * 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 React from 'react';
+import { v4 } from 'uuid';
+import { isEmpty, omitBy } from 'lodash';
+import { useEuiTheme } from '@elastic/eui';
+import { i18n } from '@kbn/i18n';
+import { FormattedMessage } from '@kbn/i18n-react';
+import type { AuthenticatedUser } from '@kbn/security-plugin/common';
+import { Message, MessageRole } from '../../common';
+import type { ChatTimelineItem } from '../components/chat/chat_timeline';
+import { RenderFunction } from '../components/render_function';
+import type { ObservabilityAIAssistantChatService } from '../types';
+
+function convertMessageToMarkdownCodeBlock(message: Message['message']) {
+  let value: object;
+
+  if (!message.name) {
+    const name = message.function_call?.name;
+    const args = message.function_call?.arguments
+      ? JSON.parse(message.function_call.arguments)
+      : undefined;
+
+    value = {
+      name,
+      args,
+    };
+  } else {
+    const content = message.content ? JSON.parse(message.content) : undefined;
+    const data = message.data ? JSON.parse(message.data) : undefined;
+    value = omitBy(
+      {
+        content,
+        data,
+      },
+      isEmpty
+    );
+  }
+
+  return `\`\`\`\n${JSON.stringify(value, null, 2)}\n\`\`\``;
+}
+
+function FunctionName({ name: functionName }: { name: string }) {
+  const { euiTheme } = useEuiTheme();
+
+  return <span style={{ fontFamily: euiTheme.font.familyCode, fontSize: 13 }}>{functionName}</span>;
+}
+
+export function getTimelineItemsfromConversation({
+  currentUser,
+  messages,
+  hasConnector,
+  chatService,
+}: {
+  currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
+  messages: Message[];
+  hasConnector: boolean;
+  chatService: ObservabilityAIAssistantChatService;
+}): ChatTimelineItem[] {
+  return [
+    {
+      id: v4(),
+      actions: { canCopy: false, canEdit: false, canGiveFeedback: false, canRegenerate: false },
+      display: { collapsed: false, hide: false },
+      currentUser,
+      loading: false,
+      role: MessageRole.User,
+      title: i18n.translate('xpack.observabilityAiAssistant.conversationStartTitle', {
+        defaultMessage: 'started a conversation',
+      }),
+    },
+    ...messages.map((message, index) => {
+      const id = v4();
+
+      let title: React.ReactNode = '';
+      let content: string | undefined;
+      let element: React.ReactNode | undefined;
+
+      const prevFunctionCall =
+        message.message.name && messages[index - 1] && messages[index - 1].message.function_call
+          ? messages[index - 1].message.function_call
+          : undefined;
+
+      const role = message.message.function_call?.trigger || message.message.role;
+
+      const actions = {
+        canCopy: false,
+        canEdit: false,
+        canGiveFeedback: false,
+        canRegenerate: false,
+      };
+
+      const display = {
+        collapsed: false,
+        hide: false,
+      };
+
+      switch (role) {
+        case MessageRole.System:
+          display.hide = true;
+          break;
+
+        case MessageRole.User:
+          actions.canCopy = true;
+          actions.canGiveFeedback = false;
+          actions.canRegenerate = false;
+
+          display.hide = false;
+
+          // User executed a function:
+          if (message.message.name && prevFunctionCall) {
+            const parsedContent = JSON.parse(message.message.content ?? 'null');
+            const isError = !!(parsedContent && 'error' in parsedContent);
+
+            title = !isError ? (
+              <FormattedMessage
+                id="xpack.observabilityAiAssistant.userExecutedFunctionEvent"
+                defaultMessage="executed the function {functionName}"
+                values={{
+                  functionName: <FunctionName name={message.message.name} />,
+                }}
+              />
+            ) : (
+              <FormattedMessage
+                id="xpack.observabilityAiAssistant.executedFunctionFailureEvent"
+                defaultMessage="failed to execute the function {functionName}"
+                values={{
+                  functionName: <FunctionName name={message.message.name} />,
+                }}
+              />
+            );
+
+            element =
+              !isError && chatService.hasRenderFunction(message.message.name) ? (
+                <RenderFunction
+                  name={message.message.name}
+                  arguments={prevFunctionCall?.arguments}
+                  response={message.message}
+                />
+              ) : undefined;
+
+            content = !element ? convertMessageToMarkdownCodeBlock(message.message) : undefined;
+
+            actions.canEdit = false;
+            display.collapsed = !isError && !element;
+          } else if (message.message.function_call) {
+            // User suggested a function
+            title = (
+              <FormattedMessage
+                id="xpack.observabilityAiAssistant.userSuggestedFunctionEvent"
+                defaultMessage="requested the function {functionName}"
+                values={{
+                  functionName: <FunctionName name={message.message.function_call.name} />,
+                }}
+              />
+            );
+
+            content = convertMessageToMarkdownCodeBlock(message.message);
+
+            actions.canEdit = hasConnector;
+            display.collapsed = true;
+          } else {
+            // is a prompt by the user
+            title = '';
+            content = message.message.content;
+
+            actions.canEdit = hasConnector;
+            display.collapsed = false;
+          }
+
+          break;
+
+        case MessageRole.Assistant:
+          actions.canRegenerate = hasConnector;
+          actions.canCopy = true;
+          actions.canGiveFeedback = true;
+          display.hide = false;
+
+          // is a function suggestion by the assistant
+          if (message.message.function_call?.name) {
+            title = (
+              <FormattedMessage
+                id="xpack.observabilityAiAssistant.suggestedFunctionEvent"
+                defaultMessage="requested the function {functionName}"
+                values={{
+                  functionName: <FunctionName name={message.message.function_call.name} />,
+                }}
+              />
+            );
+            content = convertMessageToMarkdownCodeBlock(message.message);
+
+            display.collapsed = true;
+            actions.canEdit = true;
+          } else {
+            // is an assistant response
+            title = '';
+            content = message.message.content;
+            display.collapsed = false;
+            actions.canEdit = false;
+          }
+          break;
+      }
+
+      return {
+        id,
+        role,
+        title,
+        content,
+        element,
+        actions,
+        display,
+        currentUser,
+        function_call: message.message.function_call,
+        loading: false,
+      };
+    }),
+  ];
+}

x-pack/plugins/observability_ai_assistant/public/utils/storybook_decorator.tsx

@@ -12,12 +12,34 @@ import type { AuthenticatedUser } from '@kbn/security-plugin/common';
 import { ObservabilityAIAssistantProvider } from '../context/observability_ai_assistant_provider';
 import { ObservabilityAIAssistantAPIClient } from '../api';
 import type { Message } from '../../common';
-import type { ObservabilityAIAssistantService, PendingMessage } from '../types';
+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 = {
+  chat: (options: { messages: Message[]; connectorId: string }) => new Observable<PendingMessage>(),
+  getContexts: () => [],
+  getFunctions: () => [buildFunctionElasticsearch(), buildFunctionServiceSummary()],
+  executeFunction: async (
+    name: string,
+    args: string | undefined,
+    signal: AbortSignal
+  ): Promise<{ content?: Serializable; data?: Serializable }> => ({}),
+  renderFunction: (name: string, args: string | undefined, response: {}) => (
+    <div>Hello! {name}</div>
+  ),
+  hasRenderFunction: () => true,
+};
 
 const service: ObservabilityAIAssistantService = {
   isEnabled: () => true,
-  chat: (options: { messages: Message[]; connectorId: string }) => new Observable<PendingMessage>(),
+  start: async () => {
+    return chatService;
+  },
   callApi: {} as ObservabilityAIAssistantAPIClient,
   getCurrentUser: async (): Promise<AuthenticatedUser> => ({
     username: 'user',
@@ -29,14 +51,6 @@ const service: ObservabilityAIAssistantService = {
     authentication_type: '',
     elastic_cloud_user: false,
   }),
-  getContexts: () => [],
-  getFunctions: () => [buildFunctionElasticsearch(), buildFunctionServiceSummary()],
-  executeFunction: async (
-    name: string,
-    args: string | undefined,
-    signal: AbortSignal
-  ): Promise<{ content?: Serializable; data?: Serializable }> => ({}),
-  renderFunction: (name: string, response: {}) => <div>Hello! {name}</div>,
 };
 
 export function KibanaReactStorybookDecorator(Story: ComponentType) {
@@ -54,7 +68,9 @@ export function KibanaReactStorybookDecorator(Story: ComponentType) {
       }}
     >
       <ObservabilityAIAssistantProvider value={service}>
-        <Story />
+        <ObservabilityAIAssistantChatServiceProvider value={chatService}>
+          <Story />
+        </ObservabilityAIAssistantChatServiceProvider>
       </ObservabilityAIAssistantProvider>
     </KibanaContextProvider>
   );

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

@@ -24,7 +24,6 @@ const chatRoute = createObservabilityAIAssistantServerRoute({
           name: t.string,
           description: t.string,
           parameters: t.any,
-          contexts: t.array(t.string),
         })
       ),
     }),

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

@@ -109,10 +109,36 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
   },
 });
 
+const getKnowledgeBaseStatus = createObservabilityAIAssistantServerRoute({
+  endpoint: 'GET /internal/observability_ai_assistant/functions/kb_status',
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (
+    resources
+  ): Promise<{
+    ready: boolean;
+    error?: any;
+    deployment_state?: string;
+    allocation_state?: string;
+  }> => {
+    const client = await resources.service.getClient({ request: resources.request });
+
+    if (!client) {
+      throw notImplemented();
+    }
+
+    return await client.getKnowledgeBaseStatus();
+  },
+});
+
 const setupKnowledgeBaseRoute = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/functions/setup_kb',
   options: {
     tags: ['access:ai_assistant'],
+    timeout: {
+      idleSocket: 20 * 60 * 1000, // 20 minutes
+    },
   },
   handler: async (resources): Promise<void> => {
     const client = await resources.service.getClient({ request: resources.request });
@@ -130,4 +156,5 @@ export const functionRoutes = {
   ...functionRecallRoute,
   ...functionSummariseRoute,
   ...setupKnowledgeBaseRoute,
+  ...getKnowledgeBaseStatus,
 };

x-pack/plugins/observability_ai_assistant/server/routes/types.ts

@@ -31,6 +31,9 @@ export interface ObservabilityAIAssistantRouteHandlerResources {
 
 export interface ObservabilityAIAssistantRouteCreateOptions {
   options: {
+    timeout?: {
+      idleSocket?: number;
+    };
     tags: Array<'access:ai_assistant'>;
   };
 }

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

@@ -18,9 +18,10 @@ import type {
   ChatCompletionRequestMessage,
   CreateChatCompletionRequest,
 } from 'openai';
+import pRetry from 'p-retry';
 import { v4 } from 'uuid';
 import {
-  KnowledgeBaseEntry,
+  type KnowledgeBaseEntry,
   MessageRole,
   type Conversation,
   type ConversationCreateRequest,
@@ -144,7 +145,7 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant
   }: {
     messages: Message[];
     connectorId: string;
-    functions: Array<FunctionDefinition['options']>;
+    functions: Array<Pick<FunctionDefinition['options'], 'name' | 'description' | 'parameters'>>;
   }): Promise<IncomingMessage> => {
     const messagesForOpenAI: ChatCompletionRequestMessage[] = compact(
       messages
@@ -164,9 +165,7 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant
         })
     );
 
-    const functionsForOpenAI: ChatCompletionFunctions[] = functions.map((fn) =>
-      omit(fn, 'contexts')
-    );
+    const functionsForOpenAI: ChatCompletionFunctions[] = functions;
 
     const request: Omit<CreateChatCompletionRequest, 'model'> & { model?: string } = {
       messages: messagesForOpenAI,
@@ -320,42 +319,96 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant
     }
   };
 
+  getKnowledgeBaseStatus = async () => {
+    try {
+      const modelStats = await this.dependencies.esClient.ml.getTrainedModelsStats({
+        model_id: ELSER_MODEL_ID,
+      });
+      const elserModelStats = modelStats.trained_model_stats[0];
+      const deploymentState = elserModelStats.deployment_stats?.state;
+      const allocationState = elserModelStats.deployment_stats?.allocation_status.state;
+      return {
+        ready: deploymentState === 'started' && allocationState === 'fully_allocated',
+        deployment_state: deploymentState,
+        allocation_state: allocationState,
+      };
+    } catch (error) {
+      return {
+        error: error instanceof errors.ResponseError ? error.body.error : String(error),
+        ready: false,
+      };
+    }
+  };
+
   setupKnowledgeBase = async () => {
     // if this fails, it's fine to propagate the error to the user
-    await this.dependencies.esClient.ml.putTrainedModel({
-      model_id: ELSER_MODEL_ID,
-      input: {
-        field_names: ['text_field'],
-      },
-    });
+
+    const installModel = async () => {
+      this.dependencies.logger.info('Installing ELSER model');
+      await this.dependencies.esClient.ml.putTrainedModel(
+        {
+          model_id: ELSER_MODEL_ID,
+          input: {
+            field_names: ['text_field'],
+          },
+          // @ts-expect-error
+          wait_for_completion: true,
+        },
+        { requestTimeout: '20m' }
+      );
+      this.dependencies.logger.info('Finished installing ELSER model');
+    };
+
+    try {
+      const getResponse = await this.dependencies.esClient.ml.getTrainedModels({
+        model_id: ELSER_MODEL_ID,
+        include: 'definition_status',
+      });
+
+      if (!getResponse.trained_model_configs[0]?.fully_defined) {
+        this.dependencies.logger.info('Model is not fully defined');
+        await installModel();
+      }
+    } catch (error) {
+      if (
+        error instanceof errors.ResponseError &&
+        error.body.error.type === 'resource_not_found_exception'
+      ) {
+        await installModel();
+      } else {
+        throw error;
+      }
+    }
 
     try {
       await this.dependencies.esClient.ml.startTrainedModelDeployment({
         model_id: ELSER_MODEL_ID,
+        wait_for: 'fully_allocated',
       });
-
-      const modelStats = await this.dependencies.esClient.ml.getTrainedModelsStats({
-        model_id: ELSER_MODEL_ID,
-      });
-
-      const elserModelStats = modelStats.trained_model_stats[0];
-
-      if (elserModelStats?.deployment_stats?.state !== 'started') {
-        throwKnowledgeBaseNotReady({
-          message: `Deployment has not started`,
-          deployment_stats: elserModelStats.deployment_stats,
-        });
-      }
-      return;
     } catch (error) {
-      if (
-        (error instanceof errors.ResponseError &&
-          error.body.error.type === 'resource_not_found_exception') ||
-        error.body.error.type === 'status_exception'
-      ) {
-        throwKnowledgeBaseNotReady(error.body);
+      if (error instanceof errors.ResponseError && error.body.error.type === 'status_exception') {
+        await pRetry(
+          async () => {
+            const response = await this.dependencies.esClient.ml.getTrainedModelsStats({
+              model_id: ELSER_MODEL_ID,
+            });
+
+            if (
+              response.trained_model_stats[0]?.deployment_stats?.allocation_status.state ===
+              'fully_allocated'
+            ) {
+              return Promise.resolve();
+            }
+
+            this.dependencies.logger.debug('Model is not allocated yet');
+
+            return Promise.reject(new Error('Not Ready'));
+          },
+          { factor: 1, minTimeout: 10000, maxRetryTime: 20 * 60 * 1000 }
+        );
+      } else {
+        throw error;
       }
-      throw error;
     }
   };
 }

x-pack/plugins/observability_ai_assistant/server/service/types.ts

@@ -20,7 +20,7 @@ export interface IObservabilityAIAssistantClient {
   chat: (options: {
     messages: Message[];
     connectorId: string;
-    functions: Array<FunctionDefinition['options']>;
+    functions: Array<Pick<FunctionDefinition['options'], 'name' | 'description' | 'parameters'>>;
   }) => Promise<IncomingMessage>;
   get: (conversationId: string) => Promise<Conversation>;
   find: (options?: { query?: string }) => Promise<{ conversations: Conversation[] }>;
@@ -29,6 +29,12 @@ export interface IObservabilityAIAssistantClient {
   delete: (conversationId: string) => Promise<void>;
   recall: (query: string) => Promise<{ entries: KnowledgeBaseEntry[] }>;
   summarise: (options: { entry: Omit<KnowledgeBaseEntry, '@timestamp'> }) => Promise<void>;
+  getKnowledgeBaseStatus: () => Promise<{
+    ready: boolean;
+    error?: any;
+    deployment_state?: string;
+    allocation_state?: string;
+  }>;
   setupKnowledgeBase: () => Promise<void>;
 }
 

x-pack/plugins/observability_ai_assistant/tsconfig.json

@@ -35,7 +35,9 @@
     "@kbn/io-ts-utils",
     "@kbn/std",
     "@kbn/alerting-plugin",
-    "@kbn/features-plugin"
+    "@kbn/features-plugin",
+    "@kbn/react-kibana-context-theme",
+    "@kbn/i18n-react"
   ],
   "exclude": ["target/**/*"]
 }

x-pack/test/apm_api_integration/tests/service_overview/dependencies/index.spec.ts

@@ -344,6 +344,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
         expectSnapshot(firstItem.location).toMatchInline(`
           Object {
             "agentName": "dotnet",
+            "dependencyName": "opbeans:3000",
             "environment": "production",
             "id": "5948c153c2d8989f92a9c75ef45bb845f53e200d",
             "serviceName": "opbeans-dotnet",