allows AI assistant to get info about available indices and their fields

x-pack/plugins/observability_ai_assistant/kibana.jsonc

@@ -18,7 +18,8 @@
       "security",
       "share",
       "taskManager",
-      "triggersActionsUi"
+      "triggersActionsUi",
+      "dataViews"
     ],
     "requiredBundles": ["fieldFormats", "kibanaReact", "kibanaUtils"],
     "optionalPlugins": [],

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

@@ -110,7 +110,7 @@ const esqlLanguagePlugin = () => {
 
 export function MessageText({ loading, content, onActionClick }: Props) {
   const containerClassName = css`
-    overflow-wrap: break-word;
+    overflow-wrap: anywhere;
   `;
 
   const onActionClickRef = useRef(onActionClick);

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

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

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

@@ -14,6 +14,7 @@ import { registerElasticsearchFunction } from './elasticsearch';
 import { registerKibanaFunction } from './kibana';
 import { registerLensFunction } from './lens';
 import { registerRecallFunction } from './recall';
+import { registerGetDatasetInfoFunction } from './get_dataset_info';
 import { registerSummarizationFunction } from './summarize';
 import { registerAlertsFunction } from './alerts';
 import { registerEsqlFunction } from './esql';
@@ -42,16 +43,16 @@ export async function registerFunctions({
 
       let description = dedent(
         `You are a helpful assistant for Elastic Observability. Your goal is to help the Elastic Observability users to quickly assess what is happening in their observed systems. You can help them visualise and analyze data, investigate their systems, perform root cause analysis or identify optimisation opportunities.
-        
+
         It's very important to not assume what the user is meaning. Ask them for clarification if needed.
-        
+
         If you are unsure about which function should be used and with what arguments, ask the user for clarification or confirmation.
 
         In KQL, escaping happens with double quotes, not single quotes. Some characters that need escaping are: ':()\\\
         /\". Always put a field value in double quotes. Best: service.name:\"opbeans-go\". Wrong: service.name:opbeans-go. This is very important!
 
         You can use Github-flavored Markdown in your responses. If a function returns an array, consider using a Markdown table to format the response.
-        
+
         If multiple functions are suitable, use the most specific and easy one. E.g., when the user asks to visualise APM data, use the APM functions (if available) rather than Lens.
 
         If a function call fails, do not execute it again with the same input. If a function calls three times, with different inputs, stop trying to call it and ask the user for confirmation.
@@ -67,8 +68,7 @@ export async function registerFunctions({
         Additionally, you can use the "recall" function to retrieve relevant information from the knowledge database.`;
 
         description += `Here are principles you MUST adhere to, in order:
-
-        - DO NOT make any assumptions about where and how users have stored their data.
+        - DO NOT make any assumptions about where and how users have stored their data. ALWAYS first call get_dataset_info function with empty string to get information about available indices. Once you know about available indices you MUST use this function again to get a list of available fields for specific index. If user provides an index name make sure its a valid index first before using it to retrieve the field list by calling this function with an empty string!
         `;
         registerSummarizationFunction({ service, registerFunction });
         registerRecallFunction({ service, registerFunction });
@@ -81,6 +81,7 @@ export async function registerFunctions({
       registerEsqlFunction({ service, registerFunction });
       registerKibanaFunction({ service, registerFunction, coreStart });
       registerAlertsFunction({ service, registerFunction });
+      registerGetDatasetInfoFunction({ service, registerFunction });
 
       registerContext({
         name: 'core',

x-pack/plugins/observability_ai_assistant/public/functions/lens.tsx

@@ -38,23 +38,24 @@ function Lens({
   end,
   lens,
   dataViews,
+  timeField,
 }: {
   indexPattern: string;
   xyDataLayer: XYDataLayer;
   start: string;
   end: string;
   lens: LensPublicStart;
   dataViews: DataViewsServicePublic;
+  timeField: string;
 }) {
   const formulaAsync = useAsync(() => {
     return lens.stateHelperApi();
   }, [lens]);
 
   const dataViewAsync = useAsync(() => {
     return dataViews.create({
-      id: indexPattern,
       title: indexPattern,
-      timeFieldName: '@timestamp',
+      timeFieldName: timeField,
     });
   }, [indexPattern]);
 
@@ -199,6 +200,12 @@ export function registerLensFunction({
               required: ['label', 'formula', 'format'],
             },
           },
+          timeField: {
+            type: 'string',
+            default: '@timefield',
+            description:
+              'time field to use for XY chart. Use @timefield if its available on the index.',
+          },
           breakdown: {
             type: 'object',
             additionalProperties: false,
@@ -235,15 +242,15 @@ export function registerLensFunction({
             description: 'The end of the time range, in Elasticsearch datemath',
           },
         },
-        required: ['layers', 'indexPattern', 'start', 'end'],
+        required: ['layers', 'indexPattern', 'start', 'end', 'timeField'],
       } as const,
     },
     async () => {
       return {
         content: {},
       };
     },
-    ({ arguments: { layers, indexPattern, breakdown, seriesType, start, end } }) => {
+    ({ arguments: { layers, indexPattern, breakdown, seriesType, start, end, timeField } }) => {
       const xyDataLayer = new XYDataLayer({
         data: layers.map((layer) => ({
           type: 'formula',
@@ -263,6 +270,8 @@ export function registerLensFunction({
         },
       });
 
+      if (!timeField) return;
+
       return (
         <Lens
           indexPattern={indexPattern}
@@ -271,6 +280,7 @@ export function registerLensFunction({
           end={end}
           lens={pluginsStart.lens}
           dataViews={pluginsStart.dataViews}
+          timeField={timeField}
         />
       );
     }

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

@@ -276,11 +276,84 @@ const setupKnowledgeBaseRoute = createObservabilityAIAssistantServerRoute({
   },
 });
 
+const functionGetDatasetInfoRoute = createObservabilityAIAssistantServerRoute({
+  endpoint: 'POST /internal/observability_ai_assistant/functions/get_dataset_info',
+  params: t.type({
+    body: t.type({
+      index: t.string,
+    }),
+  }),
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (
+    resources
+  ): Promise<{
+    indices: string[];
+    fields: Array<{ name: string; description: string; type: string }>;
+  }> => {
+    const esClient = (await resources.context.core).elasticsearch.client.asCurrentUser;
+
+    const savedObjectsClient = (await resources.context.core).savedObjects.getClient();
+
+    const index = resources.params.body.index;
+
+    let indices: string[] = [];
+
+    try {
+      const body = await esClient.indices.resolveIndex({
+        name: index === '' ? '*' : index,
+        expand_wildcards: 'open',
+      });
+      indices = [...body.indices.map((i) => i.name), ...body.data_streams.map((d) => d.name)];
+    } catch (e) {
+      indices = [];
+    }
+
+    if (index === '') {
+      return {
+        indices,
+        fields: [],
+      };
+    }
+
+    if (indices.length === 0) {
+      return {
+        indices,
+        fields: [],
+      };
+    }
+
+    const dataViews = await (
+      await resources.plugins.dataViews.start()
+    ).dataViewsServiceFactory(savedObjectsClient, esClient);
+
+    const fields = await dataViews.getFieldsForWildcard({
+      pattern: index,
+    });
+
+    // else get all the fields for the found dataview
+    return {
+      indices: [index],
+      fields: fields.flatMap((field) => {
+        return (field.esTypes ?? [field.type]).map((type) => {
+          return {
+            name: field.name,
+            description: field.customLabel || '',
+            type,
+          };
+        });
+      }),
+    };
+  },
+});
+
 export const functionRoutes = {
   ...functionElasticsearchRoute,
   ...functionRecallRoute,
   ...functionSummariseRoute,
   ...setupKnowledgeBaseRoute,
   ...getKnowledgeBaseStatus,
   ...functionAlertsRoute,
+  ...functionGetDatasetInfoRoute,
 };

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

@@ -17,6 +17,7 @@ import type {
   TaskManagerSetupContract,
   TaskManagerStartContract,
 } from '@kbn/task-manager-plugin/server';
+import { DataViewsServerPluginStart } from '@kbn/data-views-plugin/server';
 
 /* eslint-disable @typescript-eslint/no-empty-interface*/
 export interface ObservabilityAIAssistantPluginStart {}
@@ -32,4 +33,5 @@ export interface ObservabilityAIAssistantPluginStartDependencies {
   security: SecurityPluginStart;
   features: FeaturesPluginStart;
   taskManager: TaskManagerStartContract;
+  dataViews: DataViewsServerPluginStart;
 }